Add n_frames parameter for vibration trajectory frame count

README.md

@@ -443,11 +443,12 @@ graphrc <input_file> [options]
 > graphrc -h
 
 usage: graphrc [-h] [--version] [--cite] [--mode MODE] [--ts-frame TS_FRAME] [--frame-selection {rmsd,bookend}] [--relaxed] [--bond-tolerance BOND_TOLERANCE]
-               [--bond-threshold BOND_THRESHOLD] [--angle-threshold ANGLE_THRESHOLD] [--dihedral-threshold DIHEDRAL_THRESHOLD] [--coupled-motion-filter COUPLED_MOTION_FILTER]
-               [--coupled-proton-threshold COUPLED_PROTON_THRESHOLD] [--all] [--graph] [--method {cheminf,xtb}] [--charge CHARGE] [--multiplicity MULTIPLICITY]
-               [--distance-tolerance DISTANCE_TOLERANCE] [--independent-graphs] [--ig-flexible] [--ascii-scale ASCII_SCALE] [--show-h] [--ascii-shells ASCII_SHELLS]
-               [--save-displacement] [--displacement-scale DISPLACEMENT_SCALE] [--no-save] [--orca-path ORCA_PATH] [--debug]
-               [input]
+              [--bond-threshold BOND_THRESHOLD] [--angle-threshold ANGLE_THRESHOLD] [--dihedral-threshold DIHEDRAL_THRESHOLD]
+              [--coupled-motion-filter COUPLED_MOTION_FILTER] [--coupled-proton-threshold COUPLED_PROTON_THRESHOLD] [--all] [--graph] [--method {cheminf,xtb}]
+              [--charge CHARGE] [--multiplicity MULTIPLICITY] [--distance-tolerance DISTANCE_TOLERANCE] [--independent-graphs] [--ig-flexible]
+              [--ascii-scale ASCII_SCALE] [--show-h] [--ascii-shells ASCII_SHELLS] [--vib-frames VIB_FRAMES] [--save-displacement]
+              [--displacement-scale DISPLACEMENT_SCALE] [--no-save] [--debug]
+              [input]
 
 Internal Coordinate Analysis of Vibrational Modes.
 
@@ -490,10 +491,10 @@ graph analysis parameters:
   --distance-tolerance DISTANCE_TOLERANCE
                         Tolerance for bond formation/breaking (default: 0.2 Å)
   --independent-graphs, -ig
-                        Build molecular graphs from the displaced geometries rather than TS geometry with guided bonding (more rigorous for use
-                        with IRC or QRC displaced trajectories)
-  --ig-flexible, -igf   Apply bond-tolerance to displaced graphs (with -ig). Default: displaced graphs use stricter xyzgraph defaults for more
-                        rigorous connectivity detection
+                        Build molecular graphs from the displaced geometries rather than TS geometry with guided bonding (more rigorous for use with IRC or
+                        QRC displaced trajectories)
+  --ig-flexible, -igf   Apply bond-tolerance to displaced graphs (with -ig). Default: displaced graphs use stricter xyzgraph defaults for more rigorous
+                        connectivity detection
 
 ASCII rendering options:
   --ascii-scale ASCII_SCALE, -as ASCII_SCALE
@@ -503,11 +504,14 @@ ASCII rendering options:
                         Neighbor shells around transformation core (default: 1)
 
 output options:
+  --vib-frames VIB_FRAMES
+                        Number of frames in the generated trajectory; must be a positive multiple of 4 (invalid values fall back to the default; QM output
+                        only; default: 20)
   --save-displacement, -sd
                         Save displaced structure pair
   --displacement-scale DISPLACEMENT_SCALE, -ds DISPLACEMENT_SCALE
-                        Displacement scale for saved displaced structures (1-4 uses trajectory frames directly, ~0.2-0.8
-                        amplitude; >4 extrapolates beyond the trajectory using the normal mode vector) (default: 1)
+                        Displacement scale for saved displaced structures (1-4 uses trajectory frames directly, ~0.2-0.8 amplitude; >4 extrapolates beyond the
+                        trajectory) (default: 1)
   --no-save             Do not save trajectory to disk (keep in memory only)
 ```
 
@@ -589,6 +593,9 @@ graphrc input.xyz -sd -ds 2
 
 # Don't save trajectory to disk
 graphrc input.xyz --no-save
+
+# Set the trajectory frame count for QM output (must be a positive multiple of 4)
+graphrc calculation.out --vib-frames 40
 ```
 
 ### Complete Example

src/graphrc/api.py

@@ -23,7 +23,13 @@
 
 from . import config
 from .characterize import characterize_vib_mode
-from .convert import is_orca_output, parse_cclib_output, parse_orca_output, parse_xyz_string_to_frames
+from .convert import (
+    is_orca_output,
+    parse_cclib_output,
+    parse_orca_output,
+    parse_xyz_string_to_frames,
+    validate_n_frames,
+)
 from .core import analyze_internal_displacements, read_xyz_trajectory
 from .graph_compare import analyze_displacement_graphs
 from .utils import save_displacement_pair, setup_logging, write_trajectory_file
@@ -63,6 +69,7 @@ def collect_metadata(input_file: str, **params) -> Dict[str, Any]:
 def load_trajectory(
     input_file: str,
     mode: int = 0,
+    n_frames: int = config.DEFAULT_N_FRAMES,
     save_to_disk: bool = True,
     print_output: bool = False,
 ) -> Dict[str, Any]:
@@ -72,6 +79,9 @@ def load_trajectory(
     Args:
         input_file: Path to XYZ trajectory or QM output file
         mode: Vibrational mode index (ignored for XYZ files)
+        n_frames: Frames in the generated trajectory; must be a positive multiple
+            of 4. Only affects QM output — XYZ frame count is read from the file
+            (default 20)
         save_to_disk: Whether to save converted trajectory to disk
         print_output: Print status messages
 
@@ -82,6 +92,10 @@ def load_trajectory(
             - 'frequencies': List of frequencies (None for XYZ input)
             - 'trajectory_file': Path to trajectory file (None if not saved)
     """
+    # Validate up front so a bad value fails fast with a clear message, before any
+    # file IO. n_frames only affects QM output, but is always checked for sanity.
+    validate_n_frames(n_frames)
+
     basename = os.path.basename(input_file)
     root, ext = os.path.splitext(input_file)
 
@@ -104,15 +118,15 @@ def load_trajectory(
         try:
             if print_output:
                 print(f"\nParsing {basename} with direct ORCA output parser...")
-            frequencies, trajectory_string, mode_displacement = parse_orca_output(input_file, mode)
+            frequencies, trajectory_string, mode_displacement = parse_orca_output(input_file, mode, n_frames)
         except Exception as e:
             if print_output:
                 print(f"Direct ORCA parsing failed ({e}), falling back to cclib...")
-            frequencies, trajectory_string, mode_displacement = parse_cclib_output(input_file, mode)
+            frequencies, trajectory_string, mode_displacement = parse_cclib_output(input_file, mode, n_frames)
     else:
         if print_output:
             print(f"\nParsing {basename} with cclib...")
-        frequencies, trajectory_string, mode_displacement = parse_cclib_output(input_file, mode)
+        frequencies, trajectory_string, mode_displacement = parse_cclib_output(input_file, mode, n_frames)
 
     # Convert string to frames
     frames = parse_xyz_string_to_frames(trajectory_string)
@@ -156,6 +170,8 @@ def run_vib_analysis(
     ascii_scale: float = config.ASCII_SCALE,
     ascii_include_h: bool = config.ASCII_INCLUDE_H,
     ascii_neighbor_shells: int = config.ASCII_NEIGHBOR_SHELLS,
+    # Trajectory frame count (QM output files only; ignored for XYZ input)
+    n_frames: int = config.DEFAULT_N_FRAMES,
     # Output options
     save_trajectory: bool = config.SAVE_TRAJECTORY_DEFAULT,
     save_displacement: bool = config.SAVE_DISPLACEMENT_DEFAULT,
@@ -191,6 +207,9 @@ def run_vib_analysis(
         ascii_include_h: Include hydrogens in ASCII rendering
         ascii_neighbor_shells: Neighbor shells around transformation core
 
+        n_frames: Frames in the generated trajectory; must be a positive multiple
+            of 4. Only affects QM output — XYZ frame count is read from the file
+            (default 20)
         save_trajectory: Save converted trajectory to disk
         save_displacement: Save displaced structure pair
         displacement_scale: Displacement amplitude level (1-4)
@@ -207,6 +226,9 @@ def run_vib_analysis(
             - 'graph': Graph analysis results (if enabled)
             - 'displacement_files': Paths to saved displacement files (if enabled)
     """
+    # Fail fast on an invalid frame count before printing anything or doing work.
+    validate_n_frames(n_frames)
+
     if print_output or debug:
         print("=" * 80)
         print(" " * 32 + "GRAPHRC")
@@ -243,6 +265,7 @@ def run_vib_analysis(
     trajectory_data = load_trajectory(
         input_file,
         mode=mode,
+        n_frames=n_frames,
         save_to_disk=save_trajectory,
         print_output=False,
     )

src/graphrc/cli.py

@@ -6,6 +6,7 @@
 
 from . import __citation__, __version__, config
 from .api import run_vib_analysis
+from .convert import validate_n_frames
 
 
 def main():
@@ -139,6 +140,13 @@ def main():
 
     # Output options
     output_group = parser.add_argument_group("output options")
+    output_group.add_argument(
+        "--vib-frames",
+        type=int,
+        default=config.DEFAULT_N_FRAMES,
+        help=f"Number of frames in the generated trajectory; must be a positive multiple of 4 "
+        f"(invalid values fall back to the default; QM output only; default: {config.DEFAULT_N_FRAMES})",
+    )
     output_group.add_argument("--save-displacement", "-sd", action="store_true", help="Save displaced structure pair")
     output_group.add_argument(
         "--displacement-scale",
@@ -192,6 +200,18 @@ def main():
             )
             sys.exit(1)
 
+    # Validate --vib-frames; fall back to the default if it is not a positive
+    # multiple of 4. The API enforces this strictly; the CLI is forgiving.
+    try:
+        validate_n_frames(args.vib_frames)
+    except ValueError:
+        print(
+            f"Warning: --vib-frames must be a positive multiple of 4; "
+            f"using default ({config.DEFAULT_N_FRAMES}) instead of {args.vib_frames}.",
+            file=sys.stderr,
+        )
+        args.vib_frames = config.DEFAULT_N_FRAMES
+
     # Auto-enable -ig if -igf is used
     if args.ig_flexible and not args.independent_graphs:
         print("Warning: --ig-flexible requires -ig, automatically enabling -ig", file=sys.stderr)
@@ -224,6 +244,7 @@ def main():
             ascii_include_h=args.show_h,
             ascii_neighbor_shells=args.ascii_shells,
             # Output options
+            n_frames=args.vib_frames,
             save_trajectory=not args.no_save,
             save_displacement=args.save_displacement,
             displacement_scale=args.displacement_scale,

src/graphrc/config.py

@@ -51,3 +51,4 @@
 SAVE_TRAJECTORY_DEFAULT = True
 SAVE_DISPLACEMENT_DEFAULT = False
 DEFAULT_DISPLACEMENT_LEVEL = 1
+DEFAULT_N_FRAMES = 20

src/graphrc/convert.py

@@ -11,9 +11,34 @@
 import numpy as np
 from xyzgraph import DATA
 
+from . import config
+
 logger = logging.getLogger("graphrc")
 
 
+def validate_n_frames(n_frames: int) -> int:
+    """Return *n_frames* unchanged if it is a positive multiple of 4, else raise ValueError.
+
+    The trajectory amplitudes form a triangle wave; a multiple of 4 ensures the
+    extremes land exactly on ±1.0 and the midpoint lands exactly on 0.0.
+    """
+    if n_frames < 4 or n_frames % 4 != 0:
+        raise ValueError(f"n_frames must be a positive multiple of 4, got {n_frames}.")
+    return n_frames
+
+
+def _make_amplitudes(n_frames: int) -> np.ndarray:
+    """Generate *n_frames* amplitude samples covering one full oscillation cycle.
+
+    The cycle runs 0 → -1 → 0 → +1 → (wraps back to 0), sampled as a
+    triangle wave.  n_frames must be a multiple of 4 so that the extremes
+    land exactly on ±1.0 and the midpoint lands exactly on 0.0.
+    """
+    validate_n_frames(n_frames)
+    t = np.linspace(0, 1, n_frames, endpoint=False)
+    return np.where(t < 0.25, -4 * t, np.where(t < 0.75, 4 * t - 2, 4 - 4 * t))
+
+
 def is_orca_output(filepath: str) -> bool:
     """Return True if the file contains the ORCA banner in its header."""
     try:
@@ -26,35 +51,14 @@ def is_orca_output(filepath: str) -> bool:
     return False
 
 
-def parse_cclib_output(output_file, mode):
+def parse_cclib_output(output_file, mode, n_frames: int = config.DEFAULT_N_FRAMES):
     """
     Parse QM output with cclib and generate trajectory.
 
     Returns: (frequencies, trajectory_xyz_string).
     """
     mode = int(mode)
-    amplitudes = [
-        0.0,
-        -0.2,
-        -0.4,
-        -0.6,
-        -0.8,
-        -1.0,
-        -0.8,
-        -0.6,
-        -0.4,
-        -0.2,
-        0.0,
-        0.2,
-        0.4,
-        0.6,
-        0.8,
-        1.0,
-        0.8,
-        0.6,
-        0.4,
-        0.2,
-    ]
+    amplitudes = _make_amplitudes(n_frames)
 
     try:
         parser = cclib.io.ccopen(output_file)
@@ -90,7 +94,7 @@ def parse_cclib_output(output_file, mode):
     return freqs, trj_data, displacement
 
 
-def parse_orca_output(orca_file: str, mode: int):
+def parse_orca_output(orca_file: str, mode: int, n_frames: int = config.DEFAULT_N_FRAMES):
     """
     Parse ORCA output file directly for normal modes.
 
@@ -105,28 +109,7 @@ def parse_orca_output(orca_file: str, mode: int):
     Returns: (frequencies, trajectory_xyz_string).
     """
     mode = int(mode)
-    amplitudes = [
-        0.0,
-        -0.2,
-        -0.4,
-        -0.6,
-        -0.8,
-        -1.0,
-        -0.8,
-        -0.6,
-        -0.4,
-        -0.2,
-        0.0,
-        0.2,
-        0.4,
-        0.6,
-        0.8,
-        1.0,
-        0.8,
-        0.6,
-        0.4,
-        0.2,
-    ]
+    amplitudes = _make_amplitudes(n_frames)
 
     with open(orca_file) as f:
         lines = f.readlines()

tests/test_api.py

@@ -1,11 +1,10 @@
 """Smoke tests for graphrc API."""
 
-from pathlib import Path
+import numpy as np
+import pytest
 
 from graphrc import __citation__, __version__
-
-DATA_DIR = Path(__file__).parent.parent / "examples" / "data"
-TEST_FILE = str(DATA_DIR / "sn2.v000.xyz")
+from graphrc.convert import _make_amplitudes, validate_n_frames
 
 
 def test_version():
@@ -16,3 +15,21 @@ def test_version():
 def test_citation():
     """Citation contains package name."""
     assert "graphRC" in __citation__
+
+
+def test_make_amplitudes_default():
+    """The 20-frame default reproduces the original hardcoded amplitude list."""
+    expected = [
+        0.0, -0.2, -0.4, -0.6, -0.8, -1.0, -0.8, -0.6, -0.4, -0.2,
+        0.0, 0.2, 0.4, 0.6, 0.8, 1.0, 0.8, 0.6, 0.4, 0.2,
+    ]  # fmt: skip
+    assert np.allclose(_make_amplitudes(20), expected)
+    assert len(_make_amplitudes(40)) == 40
+
+
+def test_validate_n_frames():
+    """n_frames must be a positive multiple of 4."""
+    assert validate_n_frames(20) == 20
+    for bad in (0, 7, -4):
+        with pytest.raises(ValueError, match="multiple of 4"):
+            validate_n_frames(bad)