Examples

Ready-to-run example scripts demonstrating common workflows.

Parsing Trajectory Files

Parse different trajectory file formats (LAMMPS dump, ASE, XYZ) into a unified (N, 3) coordinate array.

 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_indices(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(24)))
94print("Subset (50 atoms) shape:", subset.shape)

Slicing-Pipeline Contact Angle

Per-frame angle via the composable TrajectoryAnalyzer with the ray-fan extractor and the slicing fitter.

 1"""Slicing-pipeline contact-angle example.
 2
 3Runs the per-frame slicing-fit pipeline (ray-fan extractor + algebraic
 4circle fitter + interface-derived wall) on a LAMMPS dump file and prints
 5the recovered mean contact angle.
 6"""
 7
 8from wetting_angle_kit.analysis import (
 9    DensityEstimator,
10    InterfaceExtractor,
11    SpaceSampling,
12    SurfaceFitter,
13    TrajectoryAnalyzer,
14    WallDetector,
15)
16from wetting_angle_kit.analysis.temporal import TemporalAggregator
17from wetting_angle_kit.parsers import LammpsDumpParser, LammpsDumpWaterFinder
18
19# --- Step 1: Define the trajectory file ---
20filename = "../../tests/trajectories/traj_spherical_drop_4k.lammpstrj"
21
22# --- Step 2: Identify the water-oxygen atoms ---
23wat_find = LammpsDumpWaterFinder(
24    filename,
25    oxygen_type=1,
26    hydrogen_type=2,
27)
28oxygen_indices = wat_find.get_water_oxygen_indices(frame_index=0)
29print("Number of water molecules:", len(oxygen_indices))
30
31# --- Step 3: Build the trajectory analyzer ---
32# Strategies: ray-fan Gaussian extractor + slicing fitter +
33# interface-derived wall + per-frame batching.
34analyzer = TrajectoryAnalyzer(
35    parser=LammpsDumpParser(filename),
36    atom_indices=oxygen_indices,
37    droplet_geometry="spherical",
38    interface_extractor=InterfaceExtractor(
39        sampling=SpaceSampling.rays(
40            delta_azimuthal=20.0,  # 20° between slicing planes
41            delta_polar=8.0,  # 8° in-plane ray step
42        ),
43        density=DensityEstimator.gaussian(),
44    ),
45    surface_fitter=SurfaceFitter.slicing(surface_filter_offset=2.0),
46    wall_detector=WallDetector.min_plus_offset(offset=0.0),
47    temporal_aggregator=TemporalAggregator(batch_size=1),
48)
49
50# --- Step 4: Run analysis on a frame range ---
51results = analyzer.analyze([1])
52print("Mean contact angle (°):", results.mean_angle)
53print("Std across batches (°):", results.std_angle)
54
55# Per-batch detail:
56batch = results.batches[0]
57print(
58    f"Frame {batch.frames[0]}: angle = {batch.angle:.2f}°, "
59    f"per-slice σ = {batch.angle_std:.2f}°, "
60    f"rms residual = {batch.rms_residual:.2f} Å"
61)

Whole-Fit Contact Angle with Bootstrap

Whole-shape sphere fit with the wall position taken from the actual substrate atoms and a bootstrap uncertainty.

 1"""Whole-shape fit contact-angle example.
 2
 3Runs the whole-fit pipeline (full-sphere Fibonacci ray fan + algebraic
 4sphere fit + wall atoms from the trajectory) on a LAMMPS dump file,
 5with 100 bootstrap resamples for the angle uncertainty.
 6"""
 7
 8from wetting_angle_kit.analysis import (
 9    DensityEstimator,
10    InterfaceExtractor,
11    SpaceSampling,
12    SurfaceFitter,
13    TrajectoryAnalyzer,
14    WallDetector,
15)
16from wetting_angle_kit.parsers import (
17    LammpsDumpParser,
18    LammpsDumpWallParser,
19    LammpsDumpWaterFinder,
20)
21
22# --- Step 1: Define the trajectory file ---
23filename = "../../tests/trajectories/traj_spherical_drop_4k.lammpstrj"
24
25# --- Step 2: Identify water-oxygen and wall-atom indices ---
26wat_find = LammpsDumpWaterFinder(filename, oxygen_type=1, hydrogen_type=2)
27oxygen_indices = wat_find.get_water_oxygen_indices(frame_index=0)
28
29# Wall parser: ``liquid_particle_types`` lists the liquid types to EXCLUDE.
30wall_parser = LammpsDumpWallParser(filename, liquid_particle_types=[1, 2])
31carbon_indices = wall_parser.parse(frame_index=0)
32
33# --- Step 3: Build the whole-fit analyzer ---
34# Strategies: full-sphere Fibonacci ray fan + sphere fit + from_atoms wall.
35analyzer = TrajectoryAnalyzer(
36    parser=LammpsDumpParser(filename),
37    atom_indices=oxygen_indices,
38    droplet_geometry="spherical",
39    interface_extractor=InterfaceExtractor(
40        sampling=SpaceSampling.rays(n_rays_sphere=400),
41        density=DensityEstimator.gaussian(density_sigma=3.0),
42    ),
43    surface_fitter=SurfaceFitter.whole(
44        surface_filter_offset=3.0,
45        bootstrap_samples=100,
46    ),
47    wall_detector=WallDetector.from_atoms(
48        wall_atom_indices=carbon_indices,
49        method="mean_top_layer",
50        top_layer_tolerance=1.0,
51    ),
52    wall_atom_indices=carbon_indices,
53)
54
55# --- Step 4: Run the analysis ---
56batch = analyzer.analyze([1]).batches[0]
57print(f"angle = {batch.angle:.2f}° ± {batch.angle_std:.2f}° (bootstrap)")
58print(f"R     = {batch.popt[3]:.2f} Å, z_wall = {batch.z_wall:.2f} Å")
59print(f"rms residual on the shell = {batch.rms_residual:.2f} Å")

Coupled-Fit Contact Angle

Coupled hyperbolic-tangent density-model fit via CoupledFit2DAnalyzer — one angle per pooled batch. The example shows both density estimators (histogram default vs Gaussian KDE).

 1"""Coupled-fit contact-angle example.
 2
 3Runs the coupled hyperbolic-tangent fit on a 2D density grid via
 4:class:`CoupledFit2DAnalyzer`. The analyzer solves interface extraction,
 5wall detection, and surface fitting together — one robust angle per pooled
 6batch.
 7
 8Two density estimators are shown:
 9
10- :meth:`DensityEstimator.binning` (the default) — top-hat histogram
11  with geometry-aware ``dV`` normalisation. Fast and exact; intrinsically
12  noisy at low per-cell counts.
13- :meth:`DensityEstimator.gaussian` — 3D Gaussian KDE on the cell
14  centres. Smooth density field with no per-cell Poisson noise; the
15  estimator of choice when running per-frame analyses or on systems
16  with low atom density per cell.
17"""
18
19from wetting_angle_kit.analysis import (
20    CoupledFit2DAnalyzer,
21    DensityEstimator,
22)
23from wetting_angle_kit.analysis.temporal import TemporalAggregator
24from wetting_angle_kit.parsers import LammpsDumpParser, LammpsDumpWaterFinder
25
26# --- Step 1: Define the trajectory file ---
27filename = "../../tests/trajectories/traj_spherical_drop_4k.lammpstrj"
28
29# --- Step 2: Identify water-oxygen atoms ---
30wat_find = LammpsDumpWaterFinder(
31    filename,
32    oxygen_type=1,
33    hydrogen_type=2,
34)
35oxygen_indices = wat_find.get_water_oxygen_indices(frame_index=0)
36print("Number of water molecules:", len(oxygen_indices))
37
38# --- Step 3: Define the grid ---
39grid_params = {
40    "xi_0": 0.0,
41    "xi_f": 70.0,
42    "dx": 2.0,
43    "zi_0": 0.0,
44    "zi_f": 70.0,
45    "dz": 2.0,
46}
47
48# --- Step 4: Pick a density estimator ---
49# Top-hat histogram on the sampling grid (default):
50estimator = DensityEstimator.binning()
51# Swap in the Gaussian KDE for smoother per-cell density. ``density_sigma``
52# is the Gaussian kernel width; 3 Å is a sensible default for
53# room-temperature water:
54# estimator = DensityEstimator.gaussian(density_sigma=2.5)
55
56# --- Step 5: Build the analyzer ---
57analyzer = CoupledFit2DAnalyzer(
58    parser=LammpsDumpParser(filename),
59    atom_indices=oxygen_indices,
60    droplet_geometry="spherical",
61    grid_params=grid_params,
62    density_estimator=estimator,
63    # Pool 10 frames per batch — the coupled fit benefits from
64    # statistics; ``batch_size=-1`` pools the entire trajectory.
65    temporal_aggregator=TemporalAggregator(batch_size=10),
66)
67
68# --- Step 6: Run analysis on a frame range ---
69# 20 frames at batch_size=10 gives two pooled batches.
70results = analyzer.analyze(range(0, 20))
71print("Mean contact angle (°):", results.mean_angle)
72print("Std across batches (°):", results.std_angle)
73
74# Per-batch detail:
75batch = results.batches[0]
76print(
77    f"Frames {batch.frames[0]}–{batch.frames[-1]}: "
78    f"angle = {batch.angle:.2f}°, "
79    f"R_eq = {batch.model_params['R_eq']:.2f} Å, "
80    f"z_wall = {batch.model_params['zi_0']:.2f} Å"
81)

Coupled-Fit 3D Contact Angle

Full 3D coupled hyperbolic-tangent fit via CoupledFit3DAnalyzer — nine-parameter model on a (xi, yi, zi) density grid. Spherical droplets only.

 1"""Coupled-fit 3D contact-angle example.
 2
 3Runs the 3D coupled hyperbolic-tangent fit on a full ``(xi, yi, zi)``
 4density grid via :class:`CoupledFit3DAnalyzer`. The nine-parameter
 5model fits the spherical-cap interface and the wall plane
 6simultaneously, recovering a single robust angle per pooled batch.
 7
 8Only spherical droplets are supported — cylindrical droplets reduce to
 9the 2D coupled fit by translational symmetry.
10"""
11
12from wetting_angle_kit.analysis import (
13    CoupledFit3DAnalyzer,
14    DensityEstimator,
15)
16from wetting_angle_kit.analysis.temporal import TemporalAggregator
17from wetting_angle_kit.parsers import LammpsDumpParser, LammpsDumpWaterFinder
18
19# --- Step 1: Define the trajectory file ---
20filename = "../../tests/trajectories/traj_spherical_drop_4k.lammpstrj"
21
22# --- Step 2: Identify water-oxygen atoms ---
23wat_find = LammpsDumpWaterFinder(
24    filename,
25    oxygen_type=1,
26    hydrogen_type=2,
27)
28oxygen_indices = wat_find.get_water_oxygen_indices(frame_index=0)
29print("Number of water molecules:", len(oxygen_indices))
30
31# --- Step 3: Define the 3D grid ---
32# xi/yi are in the droplet-centred frame; zi is in the lab frame so
33# the wall position retains physical meaning.
34grid_params = {
35    "xi_0": -30.0,
36    "xi_f": 30.0,
37    "dx": 3.2,
38    "yi_0": -30.0,
39    "yi_f": 30.0,
40    "dy": 3.2,
41    "zi_0": 0.0,
42    "zi_f": 60.0,
43    "dz": 4.0,
44}
45
46# --- Step 4: Pick a density estimator ---
47# Top-hat histogram on the 3D sampling grid (default):
48estimator = DensityEstimator.binning()
49# Swap in the Gaussian KDE for smoother per-cell density:
50# estimator = DensityEstimator.gaussian(density_sigma=3.0)
51
52# --- Step 5: Build the analyzer ---
53analyzer = CoupledFit3DAnalyzer(
54    parser=LammpsDumpParser(filename),
55    atom_indices=oxygen_indices,
56    droplet_geometry="spherical",
57    grid_params=grid_params,
58    density_estimator=estimator,
59    # Pool all frames into a single batch — the 3D density needs more
60    # atoms than the 2D one for comparable per-cell noise.
61    temporal_aggregator=TemporalAggregator(batch_size=-1),
62)
63
64# --- Step 6: Run analysis ---
65n_frames = LammpsDumpParser(filename).frame_count()
66results = analyzer.analyze(range(0, n_frames))
67print("Mean contact angle (°):", results.mean_angle)
68
69# Per-batch detail:
70batch = results.batches[0]
71print(
72    f"Frames {batch.frames[0]}–{batch.frames[-1]}: "
73    f"angle = {batch.angle:.2f}°, "
74    f"R_eq = {batch.model_params['R_eq']:.2f} Å, "
75    f"z_wall = {batch.model_params['zi_0']:.2f} Å"
76)
77print(
78    f"Droplet centre: "
79    f"xi_c = {batch.model_params['xi_c']:.2f} Å, "
80    f"yi_c = {batch.model_params['yi_c']:.2f} Å, "
81    f"zi_c = {batch.model_params['zi_c']:.2f} Å"
82)

Visualising a Per-Frame Droplet Snapshot

Pull a single slice’s interface contour off a slicing-pipeline result and render it with DropletSlicePlotter.

 1"""End-to-end example: slicing pipeline + per-frame droplet snapshot.
 2
 3Runs the slicing-fit pipeline on a LAMMPS dump file, pulls one slice's
 4interface contour + fitted circle off the result, and renders the
 5droplet snapshot with :class:`DropletSlicePlotter`.
 6"""
 7
 8from wetting_angle_kit.analysis import (
 9    DensityEstimator,
10    InterfaceExtractor,
11    SpaceSampling,
12    SurfaceFitter,
13    TrajectoryAnalyzer,
14    WallDetector,
15)
16from wetting_angle_kit.analysis.temporal import TemporalAggregator
17from wetting_angle_kit.parsers import (
18    LammpsDumpParser,
19    LammpsDumpWallParser,
20    LammpsDumpWaterFinder,
21)
22from wetting_angle_kit.visualization import DropletSlicePlotter
23
24# --- 1. Define the input trajectory ---
25filename = "../../tests/trajectories/traj_10_3_330w_nve_4k_reajust.lammpstrj"
26frame_index = 10
27
28# --- 2. Identify water-oxygen atoms ---
29wat_find = LammpsDumpWaterFinder(filename, oxygen_type=1, hydrogen_type=2)
30oxygen_indices = wat_find.get_water_oxygen_indices(frame_index=0)
31print("Number of water molecules detected:", len(oxygen_indices))
32
33# --- 3. Read atom and wall positions for the frame ---
34parser = LammpsDumpParser(filepath=filename)
35oxygen_position = parser.parse(frame_index=frame_index, indices=oxygen_indices)
36
37# Wall parser: ``liquid_particle_types`` lists what to EXCLUDE
38# (the liquid), leaving the wall atoms.
39wall_parser = LammpsDumpWallParser(filename, liquid_particle_types=[1, 2])
40wall_coords = wall_parser.parse(frame_index=frame_index)
41
42# --- 4. Run the slicing pipeline on the chosen frame ---
43analyzer = TrajectoryAnalyzer(
44    parser=LammpsDumpParser(filename),
45    atom_indices=oxygen_indices,
46    droplet_geometry="cylinder_y",
47    interface_extractor=InterfaceExtractor(
48        sampling=SpaceSampling.rays(delta_cylinder=5.0, delta_polar=8.0),
49        density=DensityEstimator.gaussian(),
50    ),
51    surface_fitter=SurfaceFitter.slicing(surface_filter_offset=2.0),
52    wall_detector=WallDetector.min_plus_offset(offset=0.0),
53    temporal_aggregator=TemporalAggregator(batch_size=1),
54)
55batch = analyzer.analyze([frame_index]).batches[0]
56print("Per-slice contact angles (°):", batch.per_slice_angles.tolist())
57
58# --- 5. Visualise one slice ---
59plotter = DropletSlicePlotter(center=True)
60slice_idx = 0  # any 0..len(slice_surfaces)-1
61
62fig = plotter.plot_surface_points(
63    oxygen_position=oxygen_position,
64    surface_data=[batch.slice_surfaces[slice_idx]],
65    popt=batch.slice_popts[slice_idx],
66    wall_coords=wall_coords,
67    alpha=float(batch.per_slice_angles[slice_idx]),
68)
69
70fig.write_html("droplet_plot.html")
71print("Plot saved as 'droplet_plot.html'")

Angle Evolution + Density Contour Plots

The two trajectory-level plotters (AngleEvolutionPlotter and DensityContourPlotter) on the same trajectory.

 1"""End-to-end example: angle evolution + density contour plots.
 2
 3Runs both the per-frame slicing pipeline and the coupled-fit
 4analyzer on the same trajectory, then renders the two trajectory-level
 5plots: the angle evolution curve (with per-batch ±σ band and running
 6mean) and the density contour with the fitted spherical cap overlaid.
 7
 8The coupled-fit analyzer is built with the default histogram
 9estimator; pass ``density_estimator=DensityEstimator.gaussian(...)``
10to render the contour over a smoothed density field instead.
11"""
12
13from wetting_angle_kit.analysis import (
14    CoupledFit2DAnalyzer,
15    DensityEstimator,
16    InterfaceExtractor,
17    SpaceSampling,
18    SurfaceFitter,
19    TrajectoryAnalyzer,
20    WallDetector,
21)
22from wetting_angle_kit.analysis.temporal import TemporalAggregator
23from wetting_angle_kit.parsers import LammpsDumpParser, LammpsDumpWaterFinder
24from wetting_angle_kit.visualization import (
25    AngleEvolutionPlotter,
26    DensityContourPlotter,
27)
28
29filename = "../../tests/trajectories/traj_spherical_drop_4k.lammpstrj"
30
31# Water-oxygen atoms.
32wat_find = LammpsDumpWaterFinder(filename, oxygen_type=1, hydrogen_type=2)
33oxygen_indices = wat_find.get_water_oxygen_indices(frame_index=0)
34
35# --- 1. Slicing pipeline → angle evolution figure ---
36slicing = TrajectoryAnalyzer(
37    parser=LammpsDumpParser(filename),
38    atom_indices=oxygen_indices,
39    droplet_geometry="spherical",
40    interface_extractor=InterfaceExtractor(
41        sampling=SpaceSampling.rays(delta_azimuthal=20.0, delta_polar=8.0),
42        density=DensityEstimator.gaussian(),
43    ),
44    surface_fitter=SurfaceFitter.slicing(surface_filter_offset=2.0),
45    wall_detector=WallDetector.min_plus_offset(offset=0.0),
46    temporal_aggregator=TemporalAggregator(batch_size=1),
47)
48slicing_results = slicing.analyze(range(0, 24))
49
50splot = AngleEvolutionPlotter(
51    slicing_results,
52    label="spherical_4k",
53    timestep=0.5,
54    time_unit="ps",
55)
56fig_evolution = splot.plot(per_frame_std=True, running_mean=True)
57fig_evolution.write_html("angle_evolution.html")
58print("Saved angle_evolution.html")
59
60# --- 2. Coupled-fit analyzer → density contour figure ---
61coupled_fit = CoupledFit2DAnalyzer(
62    parser=LammpsDumpParser(filename),
63    atom_indices=oxygen_indices,
64    droplet_geometry="spherical",
65    grid_params={
66        "xi_0": 0.0,
67        "xi_f": 70.0,
68        "dx": 2.0,
69        "zi_0": 0.0,
70        "zi_f": 70.0,
71        "dz": 2.0,
72    },
73    # density_estimator=DensityEstimator.gaussian(density_sigma=2.5),
74    temporal_aggregator=TemporalAggregator(batch_size=10),
75)
76coupled_fit_results = coupled_fit.analyze(range(0, 100))
77
78# Pick the first batch (or pass ``coupled_fit_results`` directly to
79# average the density across all batches before contouring).
80bplot = DensityContourPlotter(coupled_fit_results.batches[0], label="spherical_4k")
81fig_density = bplot.plot()
82fig_density.write_html("density_contour.html")
83print("Saved density_contour.html")