From ce7bfa73390595758eee7e875e76c5730ac147b4 Mon Sep 17 00:00:00 2001 From: Rhys Goodall Date: Sun, 1 Mar 2026 19:07:45 -0500 Subject: [PATCH 1/4] fea: add telemetry that configures basic logger. consider swapping to structlog. swap scripts to use logging over prints. --- .gitignore | 3 + .pre-commit-config.yaml | 2 +- examples/scripts/1_introduction.py | 52 +++---- examples/scripts/2_structural_optimization.py | 134 +++++++++--------- examples/scripts/3_dynamics.py | 78 +++++----- examples/scripts/4_high_level_api.py | 84 +++++------ examples/scripts/5_workflow.py | 75 +++++----- examples/scripts/6_phonons.py | 95 +++++++------ examples/scripts/7_others.py | 130 +++++++++-------- examples/scripts/8_bechmarking.py | 20 +-- pyproject.toml | 1 + tests/test_telemetry.py | 131 +++++++++++++++++ torch_sim/__init__.py | 117 ++++++++++++++- torch_sim/autobatching.py | 19 +++ torch_sim/optimizers/fire.py | 6 +- torch_sim/runners.py | 26 +++- torch_sim/telemetry.py | 98 +++++++++++++ 17 files changed, 747 insertions(+), 324 deletions(-) create mode 100644 tests/test_telemetry.py create mode 100644 torch_sim/telemetry.py diff --git a/.gitignore b/.gitignore index 01cde4042..8fe23290f 100644 --- a/.gitignore +++ b/.gitignore @@ -43,3 +43,6 @@ uv.lock # duecredit .duecredit.p + +# agents +.agents/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 63e6892b3..bae65326c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -44,4 +44,4 @@ repos: language: python types: [python] pass_filenames: false - additional_dependencies: [ty, torch, ase, mace-torch] + additional_dependencies: [ty, torch, ase, mace-torch>=0.3.15] diff --git a/examples/scripts/1_introduction.py b/examples/scripts/1_introduction.py index 97557dc63..d872d6ff0 100644 --- a/examples/scripts/1_introduction.py +++ b/examples/scripts/1_introduction.py @@ -18,8 +18,12 @@ from torch_sim.models.lennard_jones import LennardJonesModel from torch_sim.models.mace import MaceModel, MaceUrls +from torch_sim.telemetry import configure_logging, get_logger +configure_logging(log_file="1_introduction.log") +log = get_logger(name="1_introduction") + # Set up the device and data type device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.float32 @@ -28,9 +32,9 @@ # ============================================================================ # SECTION 1: Lennard-Jones Model - Simple Classical Potential # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 1: Lennard-Jones Model") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 1: Lennard-Jones Model") +log.info("=" * 70) # Create face-centered cubic (FCC) Argon # 5.26 Å is a typical lattice constant for Ar @@ -97,19 +101,19 @@ results = lj_model(state) # Print the results -print(f"Energy: {results['energy']}") -print(f"Forces shape: {results['forces'].shape}") -print(f"Stress shape: {results['stress'].shape}") -print(f"Per-atom energies shape: {results['energies'].shape}") -print(f"Per-atom stresses shape: {results['stresses'].shape}") +log.info(f"Energy: {results['energy']}") +log.info(f"Forces shape: {results['forces'].shape}") +log.info(f"Stress shape: {results['stress'].shape}") +log.info(f"Per-atom energies shape: {results['energies'].shape}") +log.info(f"Per-atom stresses shape: {results['stresses'].shape}") # ============================================================================ # SECTION 2: MACE Model - Machine Learning Potential (Batched) # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 2: MACE Model with Batched Input") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 2: MACE Model with Batched Input") +log.info("=" * 70) # Load the raw model from the downloaded model loaded_model = mace_mp( @@ -158,10 +162,10 @@ ) # You can see their shapes are as expected -print(f"Positions: {positions.shape}") -print(f"Cell: {cell.shape}") -print(f"Atomic numbers: {atomic_numbers.shape}") -print(f"System indices: {system_idx.shape}") +log.info(f"Positions: {positions.shape}") +log.info(f"Cell: {cell.shape}") +log.info(f"Atomic numbers: {atomic_numbers.shape}") +log.info(f"System indices: {system_idx.shape}") # Now we can pass them to the model results = batched_model( @@ -175,13 +179,13 @@ ) # The energy has shape (n_systems,) as the structures in a batch -print(f"Energy shape: {results['energy'].shape}") +log.info(f"Energy shape: {results['energy'].shape}") # The forces have shape (n_atoms, 3) same as positions -print(f"Forces shape: {results['forces'].shape}") +log.info(f"Forces shape: {results['forces'].shape}") # The stress has shape (n_systems, 3, 3) same as cell -print(f"Stress shape: {results['stress'].shape}") +log.info(f"Stress shape: {results['stress'].shape}") # Check if the energy, forces, and stress are the same for the Si system across batches # Each system has 64 atoms (2x2x2 supercell of 8-atom Si diamond) @@ -194,10 +198,10 @@ ) stress_diff = torch.max(torch.abs(results["stress"][0] - results["stress"][1])) -print(f"\nMax energy difference: {energy_diff}") -print(f"Max forces difference: {forces_diff}") -print(f"Max stress difference: {stress_diff}") +log.info(f"Max energy difference: {energy_diff}") +log.info(f"Max forces difference: {forces_diff}") +log.info(f"Max stress difference: {stress_diff}") -print("\n" + "=" * 70) -print("Introduction examples completed!") -print("=" * 70) +log.info("=" * 70) +log.info("Introduction examples completed!") +log.info("=" * 70) diff --git a/examples/scripts/2_structural_optimization.py b/examples/scripts/2_structural_optimization.py index 288494ab1..8650ed2c5 100644 --- a/examples/scripts/2_structural_optimization.py +++ b/examples/scripts/2_structural_optimization.py @@ -22,9 +22,13 @@ import torch_sim as ts from torch_sim.models.lennard_jones import LennardJonesModel from torch_sim.models.mace import MaceModel, MaceUrls +from torch_sim.telemetry import configure_logging, get_logger from torch_sim.units import UnitConversion +configure_logging(log_file="2_structural_optimization.log") +log = get_logger(name="2_structural_optimization") + # Set up the device and data type device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.float32 @@ -37,9 +41,9 @@ # ============================================================================ # SECTION 1: Lennard-Jones FIRE Optimization # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 1: Lennard-Jones FIRE Optimization") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 1: Lennard-Jones FIRE Optimization") +log.info("=" * 70) # Set up the random number generator generator = torch.Generator(device=device) @@ -112,21 +116,21 @@ # Run optimization for step in range(N_steps): if step % 100 == 0: - print(f"Step {step}: Potential energy: {state.energy[0].item()} eV") + log.info(f"Step {step}: Potential energy: {state.energy[0].item()} eV") state = ts.fire_step(state=state, model=lj_model, dt_max=0.01) -print(f"Initial energy: {results['energy'][0].item()} eV") -print(f"Final energy: {state.energy[0].item()} eV") -print(f"Initial max force: {torch.max(torch.abs(results['forces'][0])).item()} eV/Å") -print(f"Final max force: {torch.max(torch.abs(state.forces[0])).item()} eV/Å") +log.info(f"Initial energy: {results['energy'][0].item()} eV") +log.info(f"Final energy: {state.energy[0].item()} eV") +log.info(f"Initial max force: {torch.max(torch.abs(results['forces'][0])).item()} eV/Å") +log.info(f"Final max force: {torch.max(torch.abs(state.forces[0])).item()} eV/Å") # ============================================================================ # SECTION 2: Batched MACE FIRE Optimization (Atomic Positions Only) # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 2: Batched MACE FIRE - Positions Only") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 2: Batched MACE FIRE - Positions Only") +log.info("=" * 70) # Load MACE model loaded_model = mace_mp( @@ -151,9 +155,9 @@ atoms_list = [si_dc, cu_dc, fe_dc] -print(f"Silicon atoms: {len(si_dc)}") -print(f"Copper atoms: {len(cu_dc)}") -print(f"Iron atoms: {len(fe_dc)}") +log.info(f"Silicon atoms: {len(si_dc)}") +log.info(f"Copper atoms: {len(cu_dc)}") +log.info(f"Iron atoms: {len(fe_dc)}") # Create batched model model = MaceModel( @@ -172,23 +176,23 @@ # Initialize FIRE optimizer state = ts.fire_init(state=state, model=model, dt_start=0.005) -print("\nRunning FIRE:") +log.info("Running FIRE:") for step in range(N_steps): if step % 20 == 0: - print(f"Step {step}, Energy: {[energy.item() for energy in state.energy]}") + log.info(f"Step {step}, Energy: {[energy.item() for energy in state.energy]}") state = ts.fire_step(state=state, model=model, dt_max=0.01) -print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") -print(f"Final energies: {[energy.item() for energy in state.energy]} eV") +log.info(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") +log.info(f"Final energies: {[energy.item() for energy in state.energy]} eV") # ============================================================================ # SECTION 3: Batched MACE Gradient Descent Optimization # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 3: Batched MACE Gradient Descent") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 3: Batched MACE Gradient Descent") +log.info("=" * 70) # Reset structures with new perturbations si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) @@ -206,22 +210,22 @@ learning_rate = 0.01 state = ts.gradient_descent_init(state=state, model=model) -print("\nRunning batched gradient descent:") +log.info("Running batched gradient descent:") for step in range(N_steps): if step % 10 == 0: - print(f"Step {step}, Energy: {[res.item() for res in state.energy]} eV") + log.info(f"Step {step}, Energy: {[res.item() for res in state.energy]} eV") state = ts.gradient_descent_step(state=state, model=model, pos_lr=learning_rate) -print(f"Initial energies: {[res.item() for res in results['energy']]} eV") -print(f"Final energies: {[res.item() for res in state.energy]} eV") +log.info(f"Initial energies: {[res.item() for res in results['energy']]} eV") +log.info(f"Final energies: {[res.item() for res in state.energy]} eV") # ============================================================================ # SECTION 4: Unit Cell Filter with Gradient Descent # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 4: Unit Cell Filter with Gradient Descent") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 4: Unit Cell Filter with Gradient Descent") +log.info("=" * 70) # Recreate structures with perturbations si_dc = bulk("Si", "diamond", a=5.21, cubic=True).repeat((2, 2, 2)) @@ -252,14 +256,14 @@ scalar_pressure=0.0, ) -print("\nRunning batched unit cell gradient descent:") +log.info("Running batched unit cell gradient descent:") for step in range(N_steps): if step % 20 == 0: P1 = -torch.trace(state.stress[0]) * UnitConversion.eV_per_Ang3_to_GPa / 3 P2 = -torch.trace(state.stress[1]) * UnitConversion.eV_per_Ang3_to_GPa / 3 P3 = -torch.trace(state.stress[2]) * UnitConversion.eV_per_Ang3_to_GPa / 3 - print( + log.info( f"Step {step}, Energy: {[energy.item() for energy in state.energy]}, " f"P1={P1:.4f} GPa, P2={P2:.4f} GPa, P3={P3:.4f} GPa" ) @@ -268,16 +272,16 @@ state=state, model=model, pos_lr=pos_lr, cell_lr=cell_lr ) -print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") -print(f"Final energies: {[energy.item() for energy in state.energy]} eV") +log.info(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") +log.info(f"Final energies: {[energy.item() for energy in state.energy]} eV") # ============================================================================ # SECTION 5: Unit Cell Filter with FIRE # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 5: Unit Cell Filter with FIRE") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 5: Unit Cell Filter with FIRE") +log.info("=" * 70) # Recreate structures with perturbations si_dc = bulk("Si", "diamond", a=5.21, cubic=True).repeat((2, 2, 2)) @@ -306,30 +310,30 @@ scalar_pressure=0.0, ) -print("\nRunning batched unit cell FIRE:") +log.info("Running batched unit cell FIRE:") for step in range(N_steps): if step % 20 == 0: P1 = -torch.trace(state.stress[0]) * UnitConversion.eV_per_Ang3_to_GPa / 3 P2 = -torch.trace(state.stress[1]) * UnitConversion.eV_per_Ang3_to_GPa / 3 P3 = -torch.trace(state.stress[2]) * UnitConversion.eV_per_Ang3_to_GPa / 3 - print( + log.info( f"Step {step}, Energy: {[energy.item() for energy in state.energy]}, " f"P1={P1:.4f} GPa, P2={P2:.4f} GPa, P3={P3:.4f} GPa" ) state = ts.fire_step(state=state, model=model) -print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") -print(f"Final energies: {[energy.item() for energy in state.energy]} eV") +log.info(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") +log.info(f"Final energies: {[energy.item() for energy in state.energy]} eV") # ============================================================================ # SECTION 6: Frechet Cell Filter with FIRE # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 6: Frechet Cell Filter with FIRE") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 6: Frechet Cell Filter with FIRE") +log.info("=" * 70) # Recreate structures with perturbations si_dc = bulk("Si", "diamond", a=5.21, cubic=True).repeat((2, 2, 2)) @@ -358,22 +362,22 @@ scalar_pressure=0.0, ) -print("\nRunning batched frechet cell filter with FIRE:") +log.info("Running batched frechet cell filter with FIRE:") for step in range(N_steps): if step % 20 == 0: P1 = -torch.trace(state.stress[0]) * UnitConversion.eV_per_Ang3_to_GPa / 3 P2 = -torch.trace(state.stress[1]) * UnitConversion.eV_per_Ang3_to_GPa / 3 P3 = -torch.trace(state.stress[2]) * UnitConversion.eV_per_Ang3_to_GPa / 3 - print( + log.info( f"Step {step}, Energy: {[energy.item() for energy in state.energy]}, " f"P1={P1:.4f} GPa, P2={P2:.4f} GPa, P3={P3:.4f} GPa" ) state = ts.fire_step(state=state, model=model) -print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") -print(f"Final energies: {[energy.item() for energy in state.energy]} eV") +log.info(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") +log.info(f"Final energies: {[energy.item() for energy in state.energy]} eV") initial_pressure = [ -torch.trace(stress).item() * UnitConversion.eV_per_Ang3_to_GPa / 3 @@ -383,15 +387,15 @@ -torch.trace(stress).item() * UnitConversion.eV_per_Ang3_to_GPa / 3 for stress in state.stress ] -print(f"Initial pressure: {initial_pressure} GPa") -print(f"Final pressure: {final_pressure} GPa") +log.info(f"Initial pressure: {initial_pressure} GPa") +log.info(f"Final pressure: {final_pressure} GPa") # ============================================================================ # SECTION 7: Batched MACE L-BFGS # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 7: Batched MACE L-BFGS") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 7: Batched MACE L-BFGS") +log.info("=" * 70) # Recreate structures with perturbations si_dc = bulk("Si", "diamond", a=5.21).repeat((2, 2, 2)) @@ -409,22 +413,22 @@ results = model(state) state = ts.lbfgs_init(state=state, model=model, alpha=70.0, step_size=1.0) -print("\nRunning L-BFGS:") +log.info("Running L-BFGS:") for step in range(N_steps): if step % 20 == 0: - print(f"Step {step}, Energy: {[energy.item() for energy in state.energy]}") + log.info(f"Step {step}, Energy: {[energy.item() for energy in state.energy]}") state = ts.lbfgs_step(state=state, model=model, max_history=100) -print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") -print(f"Final energies: {[energy.item() for energy in state.energy]} eV") +log.info(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") +log.info(f"Final energies: {[energy.item() for energy in state.energy]} eV") # ============================================================================ # SECTION 8: Batched MACE BFGS # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 8: Batched MACE BFGS") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 8: Batched MACE BFGS") +log.info("=" * 70) # Recreate structures with perturbations si_dc = bulk("Si", "diamond", a=5.21).repeat((2, 2, 2)) @@ -442,16 +446,16 @@ results = model(state) state = ts.bfgs_init(state=state, model=model, alpha=70.0) -print("\nRunning BFGS:") +log.info("Running BFGS:") for step in range(N_steps): if step % 20 == 0: - print(f"Step {step}, Energy: {[energy.item() for energy in state.energy]}") + log.info(f"Step {step}, Energy: {[energy.item() for energy in state.energy]}") state = ts.bfgs_step(state=state, model=model) -print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") -print(f"Final energies: {[energy.item() for energy in state.energy]} eV") +log.info(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") +log.info(f"Final energies: {[energy.item() for energy in state.energy]} eV") -print("\n" + "=" * 70) -print("Structural optimization examples completed!") -print("=" * 70) +log.info("=" * 70) +log.info("Structural optimization examples completed!") +log.info("=" * 70) diff --git a/examples/scripts/3_dynamics.py b/examples/scripts/3_dynamics.py index 4049b77d3..281a89d65 100644 --- a/examples/scripts/3_dynamics.py +++ b/examples/scripts/3_dynamics.py @@ -22,9 +22,13 @@ import torch_sim as ts from torch_sim.models.lennard_jones import LennardJonesModel from torch_sim.models.mace import MaceModel, MaceUrls +from torch_sim.telemetry import configure_logging, get_logger from torch_sim.units import MetalUnits as Units +configure_logging(log_file="3_dynamics.log") +log = get_logger(name="3_dynamics") + # Set up the device and data type device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.float32 @@ -48,9 +52,9 @@ # ============================================================================ # SECTION 1: Lennard-Jones NVE (Microcanonical Ensemble) # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 1: Lennard-Jones NVE Simulation") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 1: Lennard-Jones NVE Simulation") +log.info("=" * 70) # Create face-centered cubic (FCC) Argon a_len = 5.26 # Lattice constant @@ -122,7 +126,7 @@ total_energy = state.energy + ts.calc_kinetic_energy( masses=state.masses, momenta=state.momenta ) - print(f"Step {step}: Total energy: {total_energy.item():.4f} eV") + log.info(f"Step {step}: Total energy: {total_energy.item():.4f} eV") # Update state using NVE integrator state = ts.nve_step(state=state, model=lj_model, dt=dt) @@ -130,15 +134,15 @@ final_total_energy = state.energy + ts.calc_kinetic_energy( masses=state.masses, momenta=state.momenta ) -print(f"Final total energy: {final_total_energy.item():.4f} eV") +log.info(f"Final total energy: {final_total_energy.item():.4f} eV") # ============================================================================ # SECTION 2: MACE NVE Simulation # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 2: MACE NVE Simulation") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 2: MACE NVE Simulation") +log.info("=" * 70) # Load MACE model loaded_model = mace_mp( @@ -183,28 +187,28 @@ state = ts.nve_init(state=state, model=mace_model, kT=kT) # Run MD simulation -print("\nStarting NVE molecular dynamics simulation...") +log.info("Starting NVE molecular dynamics simulation...") start_time = time.perf_counter() for step in range(N_steps): total_energy = state.energy + ts.calc_kinetic_energy( masses=state.masses, momenta=state.momenta, system_idx=state.system_idx ) if step % 100 == 0: - print(f"Step {step}: Total energy: {total_energy.item():.4f} eV") + log.info(f"Step {step}: Total energy: {total_energy.item():.4f} eV") state = ts.nve_step(state=state, model=mace_model, dt=dt) end_time = time.perf_counter() -print("\nSimulation complete!") -print(f"Time taken: {end_time - start_time:.2f} seconds") -print(f"Average time per step: {(end_time - start_time) / N_steps:.4f} seconds") +log.info("Simulation complete!") +log.info(f"Time taken: {end_time - start_time:.2f} seconds") +log.info(f"Average time per step: {(end_time - start_time) / N_steps:.4f} seconds") # ============================================================================ # SECTION 3: MACE NVT Langevin Simulation # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 3: MACE NVT Langevin Simulation") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 3: MACE NVT Langevin Simulation") +log.info("=" * 70) # Create diamond cubic Silicon si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) @@ -227,7 +231,7 @@ state.rng = 1 state = ts.nvt_langevin_init(model=mace_model, state=state, kT=kT) -print("\nStarting NVT Langevin simulation...") +log.info("Starting NVT Langevin simulation...") for step in range(N_steps): if step % 100 == 0: temp = ( @@ -236,22 +240,22 @@ ) / Units.temperature ) - print(f"Step {step}: Temperature: {temp.item():.4f} K") + log.info(f"Step {step}: Temperature: {temp.item():.4f} K") state = ts.nvt_langevin_step(state=state, model=mace_model, dt=dt, kT=kT, gamma=gamma) final_temp = ( ts.calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx) / Units.temperature ) -print(f"Final temperature: {final_temp.item():.4f} K") +log.info(f"Final temperature: {final_temp.item():.4f} K") # ============================================================================ # SECTION 4: MACE NVT Nose-Hoover Simulation # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 4: MACE NVT Nose-Hoover Simulation") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 4: MACE NVT Nose-Hoover Simulation") +log.info("=" * 70) # Create diamond cubic Silicon si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) @@ -266,7 +270,7 @@ state = ts.nvt_nose_hoover_init(state=state, model=mace_model, kT=kT, dt=dt) -print("\nStarting NVT Nose-Hoover simulation...") +log.info("Starting NVT Nose-Hoover simulation...") for step in range(N_steps): if step % 100 == 0: temp = ( @@ -276,7 +280,7 @@ / Units.temperature ) invariant = float(ts.nvt_nose_hoover_invariant(state, kT=kT)) - print( + log.info( f"Step {step}: Temperature: {temp.item():.4f} K, Invariant: {invariant:.4f}" ) state = ts.nvt_nose_hoover_step(state=state, model=mace_model, dt=dt, kT=kT) @@ -285,15 +289,15 @@ ts.calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx) / Units.temperature ) -print(f"Final temperature: {final_temp.item():.4f} K") +log.info(f"Final temperature: {final_temp.item():.4f} K") # ============================================================================ # SECTION 5: MACE NPT Nose-Hoover Simulation # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 5: MACE NPT Nose-Hoover Simulation") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 5: MACE NPT Nose-Hoover Simulation") +log.info("=" * 70) # Create diamond cubic Silicon si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) @@ -324,7 +328,7 @@ state=state, model=mace_model_stress, kT=kT, dt=torch.tensor(dt) ) -print("\nRunning NVT equilibration phase...") +log.info("Running NVT equilibration phase...") for step in range(N_steps_nvt): if step % 100 == 0: temp = ( @@ -336,7 +340,7 @@ invariant = float( ts.npt_nose_hoover_invariant(state, kT=kT, external_pressure=target_pressure) ) - print( + log.info( f"Step {step}: Temperature: {temp.item():.4f} K, Invariant: {invariant:.4f}" ) state = ts.npt_nose_hoover_step( @@ -352,7 +356,7 @@ state=state, model=mace_model_stress, kT=kT, dt=torch.tensor(dt) ) -print("\nRunning NPT simulation...") +log.info("Running NPT simulation...") for step in range(N_steps_npt): if step % 100 == 0: temp = ( @@ -371,7 +375,7 @@ ) pressure = float(ts.get_pressure(stress, e_kin, volume)) xx, yy, zz = torch.diag(state.current_cell[0]) - print( + log.info( f"Step {step}: Temperature: {temp.item():.4f} K, Invariant: {invariant:.4f}, " f"Pressure: {pressure:.4f} eV/ų, " f"Cell: [{xx.item():.4f}, {yy.item():.4f}, {zz.item():.4f}]" @@ -388,7 +392,7 @@ ts.calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx) / Units.temperature ) -print(f"Final temperature: {final_temp.item():.4f} K") +log.info(f"Final temperature: {final_temp.item():.4f} K") final_stress = mace_model_stress(state)["stress"] final_volume = torch.det(state.current_cell) @@ -399,8 +403,8 @@ ), final_volume, ) -print(f"Final pressure: {final_pressure.item():.4f} eV/ų") +log.info(f"Final pressure: {final_pressure.item():.4f} eV/ų") -print("\n" + "=" * 70) -print("Molecular dynamics examples completed!") -print("=" * 70) +log.info("=" * 70) +log.info("Molecular dynamics examples completed!") +log.info("=" * 70) diff --git a/examples/scripts/4_high_level_api.py b/examples/scripts/4_high_level_api.py index 8c2b7a371..565eb7c18 100644 --- a/examples/scripts/4_high_level_api.py +++ b/examples/scripts/4_high_level_api.py @@ -23,10 +23,14 @@ import torch_sim as ts from torch_sim.models.lennard_jones import LennardJonesModel from torch_sim.models.mace import MaceModel +from torch_sim.telemetry import configure_logging, get_logger from torch_sim.trajectory import TorchSimTrajectory, TrajectoryReporter from torch_sim.units import MetalUnits +configure_logging(log_file="4_high_level_api.log") +log = get_logger(name="4_high_level_api") + SMOKE_TEST = os.getenv("CI") is not None # Set device @@ -35,9 +39,9 @@ # ============================================================================ # SECTION 1: Basic Integration with Lennard-Jones # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 1: Basic Integration with Lennard-Jones") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 1: Basic Integration with Lennard-Jones") +log.info("=" * 70) lj_model = LennardJonesModel( sigma=2.0, # Å, typical for Si-Si interaction @@ -58,16 +62,16 @@ ) final_atoms = ts.io.state_to_atoms(final_state) -print(f"Final energy: {final_state.energy.item():.4f} eV") -print(f"Final atoms: {len(final_atoms)} atoms") +log.info(f"Final energy: {final_state.energy.item():.4f} eV") +log.info(f"Final atoms: {len(final_atoms)} atoms") # ============================================================================ # SECTION 2: Integration with Trajectory Logging # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 2: Integration with Trajectory Logging") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 2: Integration with Trajectory Logging") +log.info("=" * 70) trajectory_file = "tmp/lj_trajectory.h5md" @@ -109,17 +113,17 @@ ) final_atoms = traj.get_atoms(-1) -print(f"Final energy from trajectory: {final_energy:.4f} eV") -print(f"Number of kinetic energy samples: {len(kinetic_energies)}") -print(f"Number of potential energy samples: {len(potential_energies)}") +log.info(f"Final energy from trajectory: {final_energy:.4f} eV") +log.info(f"Number of kinetic energy samples: {len(kinetic_energies)}") +log.info(f"Number of potential energy samples: {len(potential_energies)}") # ============================================================================ # SECTION 3: MACE Model with High-Level API # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 3: MACE Model with High-Level API") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 3: MACE Model with High-Level API") +log.info("=" * 70) mace = mace_mp(model="small", return_raw_model=True) mace_model = MaceModel( @@ -146,15 +150,15 @@ ) final_atoms = ts.io.state_to_atoms(final_state) -print(f"Final energy: {final_state.energy.item():.4f} eV") +log.info(f"Final energy: {final_state.energy.item():.4f} eV") # ============================================================================ # SECTION 4: Batched Integration # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 4: Batched Integration") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 4: Batched Integration") +log.info("=" * 70) fe_atoms = bulk("Fe", "fcc", a=5.26, cubic=True) fe_atoms_supercell = fe_atoms.repeat([2, 2, 2]) @@ -171,16 +175,16 @@ final_atoms = ts.io.state_to_atoms(final_state) final_fe_atoms_supercell = final_atoms[3] -print(f"Number of systems: {len(final_atoms)}") -print(f"Final energies: {[e.item() for e in final_state.energy]} eV") +log.info(f"Number of systems: {len(final_atoms)}") +log.info(f"Final energies: {[e.item() for e in final_state.energy]} eV") # ============================================================================ # SECTION 5: Batched Integration with Trajectory Reporting # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 5: Batched Integration with Trajectory Reporting") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 5: Batched Integration with Trajectory Reporting") +log.info("=" * 70) systems = [si_atoms, fe_atoms, si_atoms_supercell, fe_atoms_supercell] @@ -207,15 +211,15 @@ final_energy = traj.get_array("potential_energy")[-1] final_energies_per_atom.append(final_energy / len(traj.get_atoms(-1))) -print(f"Final energies per atom: {final_energies_per_atom}") +log.info(f"Final energies per atom: {final_energies_per_atom}") # ============================================================================ # SECTION 6: Structure Optimization # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 6: Structure Optimization") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 6: Structure Optimization") +log.info("=" * 70) final_state = ts.optimize( system=systems, @@ -225,7 +229,7 @@ init_kwargs=dict(cell_filter=ts.CellFilter.unit), ) -print(f"Final optimized energies: {[e.item() for e in final_state.energy]} eV") +log.info(f"Final optimized energies: {[e.item() for e in final_state.energy]} eV") # Add perturbations rng = np.random.default_rng() @@ -236,9 +240,9 @@ # ============================================================================ # SECTION 7: Optimization with Custom Convergence Criteria # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 7: Optimization with Custom Convergence") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 7: Optimization with Custom Convergence") +log.info("=" * 70) final_state = ts.optimize( system=systems, @@ -251,15 +255,15 @@ init_kwargs=dict(cell_filter=ts.CellFilter.unit), ) -print(f"Final converged energies: {[e.item() for e in final_state.energy]} eV") +log.info(f"Final converged energies: {[e.item() for e in final_state.energy]} eV") # ============================================================================ # SECTION 8: Pymatgen Structure Support # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 8: Pymatgen Structure Support") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 8: Pymatgen Structure Support") +log.info("=" * 70) lattice = [[5.43, 0, 0], [0, 5.43, 0], [0, 0, 5.43]] species = ["Si"] * 8 @@ -285,9 +289,9 @@ ) final_structure = ts.io.state_to_structures(final_state) -print(f"Final structure type: {type(final_structure)}") -print(f"Final energy: {final_state.energy.item():.4f} eV") +log.info(f"Final structure type: {type(final_structure)}") +log.info(f"Final energy: {final_state.energy.item():.4f} eV") -print("\n" + "=" * 70) -print("High-level API examples completed!") -print("=" * 70) +log.info("=" * 70) +log.info("High-level API examples completed!") +log.info("=" * 70) diff --git a/examples/scripts/5_workflow.py b/examples/scripts/5_workflow.py index 1a9462a86..0ed65c48e 100644 --- a/examples/scripts/5_workflow.py +++ b/examples/scripts/5_workflow.py @@ -21,6 +21,11 @@ import torch_sim as ts from torch_sim.elastic import get_bravais_type from torch_sim.models.mace import MaceModel, MaceUrls +from torch_sim.telemetry import configure_logging, get_logger + + +configure_logging(log_file="5_workflow.log") +log = get_logger(name="5_workflow") # Set device @@ -30,17 +35,17 @@ ) dtype = torch.float32 -print(f"Running on device: {device}") +log.info(f"Running on device: {device}") # ============================================================================ # SECTION 1: In-Flight Autobatching Workflow # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 1: In-Flight Autobatching Workflow") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 1: In-Flight Autobatching Workflow") +log.info("=" * 70) -print("Loading MACE model...") +log.info("Loading MACE model...") mace = mace_mp(model=MaceUrls.mace_mpa_medium, return_raw_model=True) mace_model = MaceModel( model=mace, @@ -61,12 +66,12 @@ from matbench_discovery.data import DataFiles, ase_atoms_from_zip n_structures_to_relax = 100 - print(f"Loading {n_structures_to_relax:,} structures from WBM dataset...") + log.info(f"Loading {n_structures_to_relax:,} structures from WBM dataset...") ase_atoms_list = ase_atoms_from_zip( DataFiles.wbm_initial_atoms.path, limit=n_structures_to_relax ) except ImportError: - print("matbench_discovery not available, using synthetic structures...") + log.info("matbench_discovery not available, using synthetic structures...") n_structures_to_relax = 10 ase_atoms_list = [] for _ in range(n_structures_to_relax): @@ -75,7 +80,7 @@ ase_atoms_list.append(atoms) else: n_structures_to_relax = 2 - print(f"Loading {n_structures_to_relax:,} test structures...") + log.info(f"Loading {n_structures_to_relax:,} test structures...") al_atoms = bulk("Al", "hcp", a=4.05) al_atoms.positions += 0.1 * prng.normal(size=al_atoms.positions.shape) fe_atoms = bulk("Fe", "bcc", a=2.86).repeat((2, 2, 2)) @@ -105,18 +110,18 @@ batcher.load_states(fire_states) all_completed_states, convergence_tensor, state = [], None, None -print("\nStarting optimization with autobatching...") +log.info("Starting optimization with autobatching...") batch_count = 0 while (result := batcher.next_batch(state, convergence_tensor))[0] is not None: state, completed_states = result if state is None: raise RuntimeError("Expected in-flight batch state to be available") batch_count += 1 - print(f"Batch {batch_count}: Optimizing {state.n_systems} structures") + log.info(f"Batch {batch_count}: Optimizing {state.n_systems} structures") all_completed_states.extend(completed_states) if all_completed_states: - print(f"Total completed structures: {len(all_completed_states)}") + log.info(f"Total completed structures: {len(all_completed_states)}") # Run optimization steps for this batch for _step in range(10): @@ -131,18 +136,20 @@ end_time = time.perf_counter() total_time = end_time - start_time -print("\nOptimization complete!") -print(f"Total completed structures: {len(all_completed_states)}") -print(f"Total time: {total_time:.2f} seconds") -print(f"Average time per structure: {total_time / len(all_completed_states):.2f} seconds") +log.info("Optimization complete!") +log.info(f"Total completed structures: {len(all_completed_states)}") +log.info(f"Total time: {total_time:.2f} seconds") +log.info( + f"Average time per structure: {total_time / len(all_completed_states):.2f} seconds" +) # ============================================================================ # SECTION 2: Elastic Constants Calculation # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 2: Elastic Constants Calculation") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 2: Elastic Constants Calculation") +log.info("=" * 70) # Use higher precision for elastic constants dtype_elastic = torch.float64 @@ -176,31 +183,31 @@ state=state, model=model, scalar_pressure=0.0, cell_filter=ts.CellFilter.frechet ) -print("\nRelaxing structure...") +log.info("Relaxing structure...") unit_conv = ts.units.UnitConversion for step in range(300): pressure = -torch.trace(state.stress.squeeze()) / 3 * unit_conv.eV_per_Ang3_to_GPa current_fmax = torch.max(torch.abs(state.forces.squeeze())) if step % 50 == 0: - print( + log.info( f"Step {step}, Energy: {state.energy.item():.4f} eV, " f"Pressure: {pressure.item():.4f} GPa, " f"Fmax: {current_fmax.item():.4f} eV/Å" ) if current_fmax < fmax and abs(pressure) < 1e-2: - print(f"Converged at step {step}") + log.info(f"Converged at step {step}") break state = ts.fire_step(state=state, model=model) # Get bravais type bravais_type = get_bravais_type(state) -print(f"\nBravais lattice type: {bravais_type}") +log.info(f"Bravais lattice type: {bravais_type}") # Calculate elastic tensor -print("\nCalculating elastic tensor...") +log.info("Calculating elastic tensor...") elastic_tensor = ts.elastic.calculate_elastic_tensor( state=state, model=model, bravais_type=bravais_type ) @@ -214,21 +221,21 @@ ) # Print results -print("\nElastic tensor (GPa):") +log.info("Elastic tensor (GPa):") elastic_tensor_np = elastic_tensor.detach().cpu().numpy() for row in elastic_tensor_np: - print(" " + " ".join(f"{val:10.4f}" for val in row)) + log.info(" ".join(f"{val:10.4f}" for val in row)) -print("\nElastic moduli:") -print(f" Bulk modulus (GPa): {bulk_modulus:.4f}") -print(f" Shear modulus (GPa): {shear_modulus:.4f}") -print(f" Poisson's ratio: {poisson_ratio:.4f}") -print(f" Pugh's ratio (K/G): {pugh_ratio:.4f}") +log.info("Elastic moduli:") +log.info(f" Bulk modulus (GPa): {bulk_modulus:.4f}") +log.info(f" Shear modulus (GPa): {shear_modulus:.4f}") +log.info(f" Poisson's ratio: {poisson_ratio:.4f}") +log.info(f" Pugh's ratio (K/G): {pugh_ratio:.4f}") # Interpret Pugh's ratio material_type = "ductile" if pugh_ratio > 1.75 else "brittle" -print(f" Material behavior: {material_type}") +log.info(f" Material behavior: {material_type}") -print("\n" + "=" * 70) -print("Workflow examples completed!") -print("=" * 70) +log.info("=" * 70) +log.info("Workflow examples completed!") +log.info("=" * 70) diff --git a/examples/scripts/6_phonons.py b/examples/scripts/6_phonons.py index f2870ecff..9851da862 100644 --- a/examples/scripts/6_phonons.py +++ b/examples/scripts/6_phonons.py @@ -31,6 +31,11 @@ import torch_sim as ts from torch_sim.models.mace import MaceModel, MaceUrls +from torch_sim.telemetry import configure_logging, get_logger + + +configure_logging(log_file="6_phonons.log") +log = get_logger(name="6_phonons") def require_not_none[T](value: T | None, message: str) -> T: @@ -49,9 +54,9 @@ def require_not_none[T](value: T | None, message: str) -> T: # ============================================================================ # SECTION 1: Structure Relaxation for Phonons # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 1: Structure Relaxation") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 1: Structure Relaxation") +log.info("=" * 70) # Load the MACE model loaded_model = mace_mp( @@ -79,7 +84,7 @@ def require_not_none[T](value: T | None, message: str) -> T: ) # Relax the structure -print("Relaxing structure...") +log.info("Relaxing structure...") final_state = ts.optimize( system=struct, model=model, @@ -92,32 +97,32 @@ def require_not_none[T](value: T | None, message: str) -> T: ), ) -print(f"Relaxation complete. Final energy: {final_state.energy.item():.4f} eV") +log.info(f"Relaxation complete. Final energy: {final_state.energy.item():.4f} eV") # ============================================================================ # SECTION 2: Phonon DOS Calculation # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 2: Phonon DOS Calculation") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 2: Phonon DOS Calculation") +log.info("=" * 70) # Convert state to Phonopy atoms atoms = ts.io.state_to_phonopy(final_state)[0] ph = Phonopy(atoms, supercell_matrix) # Generate displaced supercells for force constant calculation -print(f"Generating displacements (distance = {displ} Å)...") +log.info(f"Generating displacements (distance = {displ} Å)...") ph.generate_displacements(distance=displ) supercells = ph.supercells_with_displacements if supercells is None: raise ValueError("No supercells generated - check Phonopy settings") -print(f"Number of displaced supercells: {len(supercells)}") +log.info(f"Number of displaced supercells: {len(supercells)}") # Convert PhonopyAtoms to batched state for efficient computation -print("Calculating forces for displaced structures (batched)...") +log.info("Calculating forces for displaced structures (batched)...") state = ts.io.phonopy_to_state(supercells, device, dtype) results = model(state) @@ -131,12 +136,12 @@ def require_not_none[T](value: T | None, message: str) -> T: start_idx = end_idx # Produce force constants from forces -print("Producing force constants...") +log.info("Producing force constants...") ph.forces = force_sets ph.produce_force_constants() # Calculate phonon DOS -print(f"Calculating phonon DOS with mesh {mesh}...") +log.info(f"Calculating phonon DOS with mesh {mesh}...") ph.run_mesh(mesh) ph.run_total_dos() @@ -147,23 +152,23 @@ def require_not_none[T](value: T | None, message: str) -> T: if freq_points is None or dos_values is None: raise RuntimeError("Phonopy total_dos has missing frequency_points or dos arrays") -print("\nPhonon DOS calculated:") -print(f" Frequency range: {freq_points.min():.3f} to {freq_points.max():.3f} THz") -print(f" DOS values range: {dos_values.min():.6f} to {dos_values.max():.6f}") +log.info("Phonon DOS calculated:") +log.info(f" Frequency range: {freq_points.min():.3f} to {freq_points.max():.3f} THz") +log.info(f" DOS values range: {dos_values.min():.6f} to {dos_values.max():.6f}") # Check for imaginary modes (negative frequencies) if freq_points.min() < -0.1: - print(" WARNING: Imaginary modes detected (freq < -0.1 THz)") + log.info(" WARNING: Imaginary modes detected (freq < -0.1 THz)") else: - print(" No imaginary modes detected (all frequencies > -0.1 THz)") + log.info(" No imaginary modes detected (all frequencies > -0.1 THz)") # ============================================================================ # SECTION 3: Phonon Band Structure Calculation # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 3: Phonon Band Structure Calculation") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 3: Phonon Band Structure Calculation") +log.info("=" * 70) try: import seekpath @@ -179,7 +184,7 @@ def require_not_none[T](value: T | None, message: str) -> T: ) # Get high-symmetry path using seekpath - print("Finding high-symmetry path...") + log.info("Finding high-symmetry path...") seekpath_data = seekpath.get_path( (ase_atoms.cell, ase_atoms.get_scaled_positions(), ase_atoms.numbers) ) @@ -195,21 +200,21 @@ def require_not_none[T](value: T | None, message: str) -> T: n_points = 51 # Points per segment q_pts, connections = get_band_qpoints_and_path_connections(path, npoints=n_points) - print(f"Calculating phonon band structure ({len(q_pts)} q-points)...") + log.info(f"Calculating phonon band structure ({len(q_pts)} q-points)...") ph.run_band_structure(q_pts, path_connections=connections) # Get band structure data bands_dict = ph.get_band_structure_dict() - print("\nPhonon band structure calculated:") - print(f" Number of paths: {len(bands_dict['frequencies'])}") - print(f" Number of bands: {len(bands_dict['frequencies'][0][0])}") + log.info("Phonon band structure calculated:") + log.info(f" Number of paths: {len(bands_dict['frequencies'])}") + log.info(f" Number of bands: {len(bands_dict['frequencies'][0][0])}") # Visualize if not in CI mode if not SMOKE_TEST: try: import pymatviz as pmv - print("\nGenerating phonon DOS plot...") + log.info("Generating phonon DOS plot...") total_dos = require_not_none( ph.total_dos, "Phonopy total_dos not available for plotting" ) @@ -224,7 +229,7 @@ def require_not_none[T](value: T | None, message: str) -> T: ) fig_dos.show() - print("Generating phonon band structure plot...") + log.info("Generating phonon band structure plot...") ph.auto_band_structure(plot=False) band_structure = require_not_none( ph.band_structure, "Phonopy band_structure not available for plotting" @@ -243,28 +248,28 @@ def require_not_none[T](value: T | None, message: str) -> T: fig_bands.show() except ImportError: - print("pymatviz not available, skipping visualization") + log.info("pymatviz not available, skipping visualization") else: - print("Skipping visualization in CI mode") + log.info("Skipping visualization in CI mode") except ImportError as e: - print(f"Skipping band structure calculation: {e}") - print("Install seekpath for band structure calculations") + log.info(f"Skipping band structure calculation: {e}") + log.info("Install seekpath for band structure calculations") # ============================================================================ # SECTION 4: Summary # ============================================================================ -print("\n" + "=" * 70) -print("Summary") -print("=" * 70) -print("Structure: Silicon (diamond)") -print("Supercell: 2x2x2") -print(f"Number of displaced structures: {len(supercells)}") -print("Batched force calculation: Yes") -print("Phonon DOS calculated: Yes") -print(f"Frequency range: {freq_points.min():.3f} to {freq_points.max():.3f} THz") - -print("\n" + "=" * 70) -print("Phonon calculation examples completed!") -print("=" * 70) +log.info("=" * 70) +log.info("Summary") +log.info("=" * 70) +log.info("Structure: Silicon (diamond)") +log.info("Supercell: 2x2x2") +log.info(f"Number of displaced structures: {len(supercells)}") +log.info("Batched force calculation: Yes") +log.info("Phonon DOS calculated: Yes") +log.info(f"Frequency range: {freq_points.min():.3f} to {freq_points.max():.3f} THz") + +log.info("=" * 70) +log.info("Phonon calculation examples completed!") +log.info("=" * 70) diff --git a/examples/scripts/7_others.py b/examples/scripts/7_others.py index e3cec1bfa..cb5ddfd54 100644 --- a/examples/scripts/7_others.py +++ b/examples/scripts/7_others.py @@ -24,17 +24,21 @@ from torch_sim.models.lennard_jones import LennardJonesModel from torch_sim.neighbors import torch_nl_linked_cell, torch_nl_n2 from torch_sim.properties.correlations import VelocityAutoCorrelation +from torch_sim.telemetry import configure_logging, get_logger from torch_sim.units import MetalUnits as Units +configure_logging(log_file="7_others.log") +log = get_logger(name="7_others") + SMOKE_TEST = os.getenv("CI") is not None # ============================================================================ # SECTION 1: Batched Neighbor List Calculations # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 1: Batched Neighbor List Calculations") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 1: Batched Neighbor List Calculations") +log.info("=" * 70) # Create multiple atomic systems atoms_list = [ @@ -55,44 +59,44 @@ # Ensure pbc has the correct shape [n_systems, 3] pbc_tensor = torch.tensor(pbc).repeat(state.n_systems, 1) -print(f"\nBatched system with {state.n_systems} structures:") +log.info(f"Batched system with {state.n_systems} structures:") for i, atoms in enumerate(atoms_list): - print(f" Structure {i}: {atoms.get_chemical_formula()} ({len(atoms)} atoms)") + log.info(f" Structure {i}: {atoms.get_chemical_formula()} ({len(atoms)} atoms)") # Method 1: Linked cell neighbor list (efficient for large systems) -print("\nCalculating neighbor lists with linked cell method...") +log.info("Calculating neighbor lists with linked cell method...") mapping, mapping_system, shifts_idx = torch_nl_linked_cell( pos, cell, pbc_tensor, cutoff, system_idx, self_interaction ) cell_shifts = transforms.compute_cell_shifts(cell, shifts_idx, mapping_system) dds = transforms.compute_distances_with_cell_shifts(pos, mapping, cell_shifts) -print("Linked cell results:") -print(f" Mapping shape: {mapping.shape}") -print(f" Mapping system shape: {mapping_system.shape}") -print(f" Shifts idx shape: {shifts_idx.shape}") -print(f" Cell shifts shape: {cell_shifts.shape}") -print(f" Distances shape: {dds.shape}") -print(f" Total neighbor pairs: {mapping.shape[1]}") +log.info("Linked cell results:") +log.info(f" Mapping shape: {mapping.shape}") +log.info(f" Mapping system shape: {mapping_system.shape}") +log.info(f" Shifts idx shape: {shifts_idx.shape}") +log.info(f" Cell shifts shape: {cell_shifts.shape}") +log.info(f" Distances shape: {dds.shape}") +log.info(f" Total neighbor pairs: {mapping.shape[1]}") # Method 2: N^2 neighbor list (simple but slower) -print("\nCalculating neighbor lists with N^2 method...") +log.info("Calculating neighbor lists with N^2 method...") mapping_n2, mapping_system_n2, shifts_idx_n2 = torch_nl_n2( pos, cell, pbc_tensor, cutoff, system_idx, self_interaction ) cell_shifts_n2 = transforms.compute_cell_shifts(cell, shifts_idx_n2, mapping_system_n2) dds_n2 = transforms.compute_distances_with_cell_shifts(pos, mapping_n2, cell_shifts_n2) -print("N^2 method results:") -print(f" Mapping shape: {mapping_n2.shape}") -print(f" Total neighbor pairs: {mapping_n2.shape[1]}") +log.info("N^2 method results:") +log.info(f" Mapping shape: {mapping_n2.shape}") +log.info(f" Total neighbor pairs: {mapping_n2.shape[1]}") # Verify consistency if mapping.shape[1] == mapping_n2.shape[1]: - print("\n✓ Both methods found the same number of neighbors") + log.info("✓ Both methods found the same number of neighbors") else: - print( - f"\n⚠ Different neighbor counts: " + log.info( + f"⚠ Different neighbor counts: " f"linked cell={mapping.shape[1]}, N^2={mapping_n2.shape[1]}" ) @@ -100,9 +104,9 @@ # ============================================================================ # SECTION 2: Velocity Autocorrelation Function (VACF) # ============================================================================ -print("\n" + "=" * 70) -print("SECTION 2: Velocity Autocorrelation Function") -print("=" * 70) +log.info("=" * 70) +log.info("SECTION 2: Velocity Autocorrelation Function") +log.info("=" * 70) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.float64 @@ -112,10 +116,10 @@ atoms = atoms.repeat((3, 3, 3)) temperature = 50.0 # Kelvin -print("\nSystem: Solid Argon") -print(f" Temperature: {temperature} K") -print(f" Number of atoms: {len(atoms)}") -print(f" Cell size: {atoms.cell.cellpar()[:3]}") +log.info("System: Solid Argon") +log.info(f" Temperature: {temperature} K") +log.info(f" Number of atoms: {len(atoms)}") +log.info(f" Cell size: {atoms.cell.cellpar()[:3]}") # Initialize velocities with Maxwell-Boltzmann distribution MaxwellBoltzmannDistribution(atoms, temperature_K=temperature) @@ -135,10 +139,10 @@ compute_forces=True, ) -print("\nLennard-Jones parameters:") -print(f" Epsilon: {epsilon} eV") -print(f" Sigma: {sigma} Å") -print(f" Cutoff: {cutoff:.2f} Å") +log.info("Lennard-Jones parameters:") +log.info(f" Epsilon: {epsilon} eV") +log.info(f" Sigma: {sigma} Å") +log.info(f" Cutoff: {cutoff:.2f} Å") # Simulation parameters timestep = 0.001 # ps (1 fs) @@ -169,9 +173,9 @@ # Run simulation num_steps = 100 if SMOKE_TEST else 15000 -print(f"\nRunning NVE simulation for {num_steps} steps...") -print(f"VACF window size: {window_size} samples") -print(f"Correlation sampling interval: {correlation_dt} steps") +log.info(f"Running NVE simulation for {num_steps} steps...") +log.info(f"VACF window size: {window_size} samples") +log.info(f"Correlation sampling interval: {correlation_dt} steps") for step in range(num_steps): state = ts.nve_step(state=state, model=lj_model, dt=dt) @@ -181,7 +185,7 @@ total_energy = state.energy + ts.calc_kinetic_energy( masses=state.masses, momenta=state.momenta ) - print(f" Step {step}: Total energy = {total_energy.item():.4f} eV") + log.info(f" Step {step}: Total energy = {total_energy.item():.4f} eV") reporter.close() @@ -191,10 +195,10 @@ if vacf_calc.vacf is not None: vacf_data = vacf_calc.vacf.detach().cpu().numpy() - print("\nVACF calculation complete:") - print(f" Number of windows averaged: {vacf_calc._window_count}") # noqa: SLF001 - print(f" VACF at t=0: {vacf_data[0]:.4f}") - print(f" VACF decay at t_max: {vacf_data[-1]:.4f}") + log.info("VACF calculation complete:") + log.info(f" Number of windows averaged: {vacf_calc._window_count}") # noqa: SLF001 + log.info(f" VACF at t=0: {vacf_data[0]:.4f}") + log.info(f" VACF decay at t_max: {vacf_data[-1]:.4f}") # Plot VACF if not in CI mode if not SMOKE_TEST: @@ -209,34 +213,34 @@ plt.grid(visible=True, alpha=0.3) plt.tight_layout() plt.savefig("tmp/vacf_example.png", dpi=150) - print("\n✓ VACF plot saved to tmp/vacf_example.png") + log.info("✓ VACF plot saved to tmp/vacf_example.png") else: - print("\nSkipping plot generation in CI mode") + log.info("Skipping plot generation in CI mode") else: - print("\nWarning: VACF data not available") + log.info("Warning: VACF data not available") # ============================================================================ # SECTION 3: Summary # ============================================================================ -print("\n" + "=" * 70) -print("Summary") -print("=" * 70) -print("\nDemonstrated features:") -print(" 1. Batched neighbor list calculations") -print(" - Linked cell method (efficient)") -print(" - N^2 method (simple)") -print(" 2. Velocity autocorrelation function (VACF)") -print(" - NVE molecular dynamics") -print(" - Running average over time windows") -print(" - Normalized correlation decay") - -print("\nKey capabilities:") -print(" - Efficient batched computations") -print(" - Multiple neighbor list algorithms") -print(" - Advanced property calculations during MD") -print(" - Trajectory analysis and correlation functions") - -print("\n" + "=" * 70) -print("Miscellaneous examples completed!") -print("=" * 70) +log.info("=" * 70) +log.info("Summary") +log.info("=" * 70) +log.info("Demonstrated features:") +log.info(" 1. Batched neighbor list calculations") +log.info(" - Linked cell method (efficient)") +log.info(" - N^2 method (simple)") +log.info(" 2. Velocity autocorrelation function (VACF)") +log.info(" - NVE molecular dynamics") +log.info(" - Running average over time windows") +log.info(" - Normalized correlation decay") + +log.info("Key capabilities:") +log.info(" - Efficient batched computations") +log.info(" - Multiple neighbor list algorithms") +log.info(" - Advanced property calculations during MD") +log.info(" - Trajectory analysis and correlation functions") + +log.info("=" * 70) +log.info("Miscellaneous examples completed!") +log.info("=" * 70) diff --git a/examples/scripts/8_bechmarking.py b/examples/scripts/8_bechmarking.py index 10d75751e..d7cc5a1d2 100644 --- a/examples/scripts/8_bechmarking.py +++ b/examples/scripts/8_bechmarking.py @@ -18,8 +18,12 @@ import torch_sim as ts from torch_sim.models.mace import MaceModel, MaceUrls +from torch_sim.telemetry import configure_logging, get_logger +configure_logging(log_file="8_bechmarking.log") +log = get_logger(name="8_bechmarking") + SMOKE_TEST = os.getenv("CI") is not None device = torch.device( @@ -84,7 +88,7 @@ def run_torchsim_static( torch.cuda.synchronize() elapsed = time.perf_counter() - t0 times.append(elapsed) - print(f" n={n} static_time={elapsed:.6f}s") + log.info(f" n={n} static_time={elapsed:.6f}s") return times @@ -126,7 +130,7 @@ def run_torchsim_relax( torch.cuda.synchronize() elapsed = time.perf_counter() - t0 times.append(elapsed) - print(f" n={n} relax_{RELAX_STEPS}_time={elapsed:.6f}s") + log.info(f" n={n} relax_{RELAX_STEPS}_time={elapsed:.6f}s") return times @@ -160,7 +164,7 @@ def run_torchsim_nve( torch.cuda.synchronize() elapsed = time.perf_counter() - t0 times.append(elapsed) - print(f" n={n} nve_time={elapsed:.6f}s") + log.info(f" n={n} nve_time={elapsed:.6f}s") return times @@ -194,7 +198,7 @@ def run_torchsim_nvt( torch.cuda.synchronize() elapsed = time.perf_counter() - t0 times.append(elapsed) - print(f" n={n} nvt_time={elapsed:.6f}s") + log.info(f" n={n} nvt_time={elapsed:.6f}s") return times @@ -206,14 +210,14 @@ def run_torchsim_nvt( model = load_mace_model(device) # Run all benchmarks -print("=== Static benchmark ===") +log.info("=== Static benchmark ===") static_times = run_torchsim_static(N_STRUCTURES_STATIC, base_structure, model, device) -print("\n=== Relax benchmark ===") +log.info("=== Relax benchmark ===") relax_times = run_torchsim_relax(N_STRUCTURES_RELAX, base_structure, model, device) -print("\n=== NVE benchmark ===") +log.info("=== NVE benchmark ===") nve_times = run_torchsim_nve(N_STRUCTURES_NVE, base_structure, model, device) -print("\n=== NVT benchmark ===") +log.info("=== NVT benchmark ===") nvt_times = run_torchsim_nvt(N_STRUCTURES_NVT, base_structure, model, device) diff --git a/pyproject.toml b/pyproject.toml index bcbc65ce3..baf7e9daa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -103,6 +103,7 @@ ignore = [ "EM102", # Exception must not use an f-string literal, assign to variable first "ERA001", # Found commented-out code "FIX002", # Line contains TODO, consider resolving the issue + "G004", # Logging statement uses f-string -> TODO refactor for performance "N803", # Variable name should be lowercase "N806", # Uppercase letters in variable names "PLC0415", # import should be at the top-level of a file diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py new file mode 100644 index 000000000..3ac9a4f64 --- /dev/null +++ b/tests/test_telemetry.py @@ -0,0 +1,131 @@ +"""Tests for torch_sim.telemetry.""" + +import json +import logging +from pathlib import Path + +import pytest + +from torch_sim import telemetry + + +@pytest.fixture(autouse=True) +def reset_logging(tmp_path: Path): + """Restore root logger state and telemetry._configured after each test.""" + root = logging.getLogger() + original_handlers = root.handlers[:] + original_level = root.level + original_configured = telemetry._configured # noqa: SLF001 + + yield tmp_path + + root.handlers.clear() + root.handlers.extend(original_handlers) + root.setLevel(original_level) + telemetry._configured = original_configured # noqa: SLF001 + + +def test_configure_logging_adds_two_handlers(tmp_path: Path): + log_file = tmp_path / "test.log" + telemetry.configure_logging(log_level="DEBUG", log_file=log_file) + + root = logging.getLogger() + assert len(root.handlers) == 2 + + +def test_configure_logging_creates_log_file(tmp_path: Path): + log_file = tmp_path / "test.log" + telemetry.configure_logging(log_file=log_file) + + logging.getLogger("test").info("hello") + assert log_file.exists() + + +def test_file_handler_writes_valid_json(tmp_path: Path): + log_file = tmp_path / "test.log" + telemetry.configure_logging(log_level="DEBUG", log_file=log_file) + + logging.getLogger("mymodule").info("structured message", extra={}) + logging.getLogger().handlers[0].flush() + + lines = log_file.read_text().splitlines() + assert lines, "log file is empty" + record = json.loads(lines[0]) + assert record["level"] == "INFO" + assert record["message"] == "structured message" + assert "time" in record + assert "name" in record + + +def test_json_record_includes_exception(tmp_path: Path): + log_file = tmp_path / "test.log" + telemetry.configure_logging(log_level="DEBUG", log_file=log_file) + + try: + raise ValueError("boom") # noqa: TRY301 + except ValueError: + logging.getLogger("err").exception("caught") + + logging.getLogger().handlers[0].flush() + record = json.loads(log_file.read_text().splitlines()[0]) + assert "exception" in record + assert "ValueError" in record["exception"] + + +def test_console_handler_respects_log_level( + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +): + log_file = tmp_path / "test.log" + telemetry.configure_logging(log_level="WARNING", log_file=log_file) + + logging.getLogger("quiet").debug("should not appear") + logging.getLogger("quiet").warning("should appear") + + captured = capsys.readouterr() + assert "should not appear" not in captured.out + assert "should appear" in captured.out + + +def test_file_handler_captures_debug_regardless_of_console_level(tmp_path: Path): + log_file = tmp_path / "test.log" + telemetry.configure_logging(log_level="WARNING", log_file=log_file) + + logging.getLogger("verbose").debug("debug only in file") + logging.getLogger().handlers[0].flush() + + record = json.loads(log_file.read_text().splitlines()[0]) + assert record["message"] == "debug only in file" + + +def test_configure_logging_is_idempotent(tmp_path: Path): + log_file = tmp_path / "test.log" + telemetry.configure_logging(log_level="INFO", log_file=log_file) + telemetry.configure_logging(log_level="INFO", log_file=log_file) + + assert len(logging.getLogger().handlers) == 2 + + +def test_get_logger_returns_named_logger(tmp_path: Path): + log_file = tmp_path / "test.log" + telemetry.configure_logging(log_file=log_file) + + logger = telemetry.get_logger("mymodule") + assert logger.name == "mymodule" + + +def test_get_logger_auto_configures(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(telemetry, "_configured", False) + monkeypatch.setattr(telemetry, "_DEFAULT_LOG_FILE", tmp_path / "auto.log") + + logger = telemetry.get_logger("auto") + assert telemetry._configured # noqa: SLF001 + assert isinstance(logger, logging.Logger) + + +def test_get_logger_none_returns_root(tmp_path: Path): + log_file = tmp_path / "test.log" + telemetry.configure_logging(log_file=log_file) + + logger = telemetry.get_logger() + assert logger is logging.getLogger() diff --git a/torch_sim/__init__.py b/torch_sim/__init__.py index fc7e52666..d9536754f 100644 --- a/torch_sim/__init__.py +++ b/torch_sim/__init__.py @@ -1,9 +1,8 @@ """TorchSim package base module.""" -# ruff: noqa: F401 import os -from datetime import datetime from importlib.metadata import version +from pathlib import Path import torch_sim._duecredit from torch_sim import ( @@ -98,7 +97,119 @@ PKG_DIR = os.path.dirname(__file__) ROOT = os.path.dirname(PKG_DIR) SCRIPTS_DIR = f"{ROOT}/examples" +TORCH_SIM_CONFIG_DIR = Path.home() / ".torchsim" +TORCH_SIM_CONFIG_DIR.mkdir(parents=True, exist_ok=True) __version__ = version("torch-sim-atomistic") +__all__ = [ + "CELL_FILTER_REGISTRY", + "INTEGRATOR_REGISTRY", + "INTEGRATOR_REGISTRY", + "OPTIM_REGISTRY", + "OPTIM_REGISTRY", + "BFGSState", + "BinningAutoBatcher", + "BinningAutoBatcher", + "CellBFGSState", + "CellFilter", + "CellFireState", + "CellLBFGSState", + "CellOptimState", + "CorrelationCalculator", + "FireState", + "InFlightAutoBatcher", + "InFlightAutoBatcher", + "Integrator", + "Integrator", + "LBFGSState", + "NPTLangevinState", + "NPTLangevinState", + "NPTNoseHooverState", + "NPTNoseHooverState", + "NVTNoseHooverState", + "NVTNoseHooverState", + "OptimState", + "Optimizer", + "SimState", + "SwapMCState", + "SwapMCState", + "TorchSimTrajectory", + "TrajectoryReporter", + "autobatching", + "bfgs_init", + "bfgs_step", + "calc_kT", + "calc_kinetic_energy", + "calc_temperature", + "concatenate_states", + "constraints", + "elastic", + "fire_init", + "fire_step", + "generate_energy_convergence_fn", + "generate_force_convergence_fn", + "get_cell_filter", + "get_pressure", + "gradient_descent_init", + "gradient_descent_step", + "initialize_state", + "integrate", + "io", + "lbfgs_init", + "lbfgs_step", + "math", + "models", + "monte_carlo", + "neighbors", + "npt_crescale_anisotropic_step", + "npt_crescale_anisotropic_step", + "npt_crescale_init", + "npt_crescale_init", + "npt_crescale_isotropic_step", + "npt_crescale_isotropic_step", + "npt_langevin_init", + "npt_langevin_init", + "npt_langevin_step", + "npt_langevin_step", + "npt_nose_hoover_init", + "npt_nose_hoover_init", + "npt_nose_hoover_invariant", + "npt_nose_hoover_invariant", + "npt_nose_hoover_step", + "npt_nose_hoover_step", + "nve_init", + "nve_init", + "nve_step", + "nve_step", + "nvt_langevin_init", + "nvt_langevin_init", + "nvt_langevin_step", + "nvt_langevin_step", + "nvt_nose_hoover_init", + "nvt_nose_hoover_init", + "nvt_nose_hoover_invariant", + "nvt_nose_hoover_invariant", + "nvt_nose_hoover_step", + "nvt_nose_hoover_step", + "nvt_vrescale_init", + "nvt_vrescale_init", + "nvt_vrescale_step", + "nvt_vrescale_step", + "optimize", + "optimizers", + "quantities", + "runners", + "state", + "static", + "swap_mc_init", + "swap_mc_init", + "swap_mc_step", + "swap_mc_step", + "system_wise_max_force", + "trajectory", + "transforms", + "units", +] -import torch_sim._citations # noqa: E402 + +import torch_sim._citations # noqa: E402, F401 diff --git a/torch_sim/autobatching.py b/torch_sim/autobatching.py index 36902eb7f..6925e4298 100644 --- a/torch_sim/autobatching.py +++ b/torch_sim/autobatching.py @@ -20,6 +20,7 @@ model architectures and GPU configurations. """ +import logging from collections.abc import Callable, Iterator, Sequence from itertools import chain from typing import Any, get_args @@ -32,6 +33,9 @@ from torch_sim.typing import MemoryScaling +logger = logging.getLogger(__name__) + + def to_constant_volume_bins( # noqa: C901, PLR0915 items: dict[int, float] | list[Any], max_volume: float, @@ -606,6 +610,7 @@ def load_states(self, states: T | Sequence[T]) -> float: oom_error_message=self.oom_error_message, ) self.max_memory_scaler = self.max_memory_scaler * self.max_memory_padding + logger.debug("Estimated max memory scaler: %.3g", self.max_memory_scaler) # verify that no systems are too large max_metric_value = max(self.memory_scalers) @@ -627,6 +632,12 @@ def load_states(self, states: T | Sequence[T]) -> float: self.batched_states = [[batched[index_bin]] for index_bin in self.index_bins] self.current_state_bin = 0 + logger.info( + "BinningAutoBatcher: %d systems → %d batch(es), max_memory_scaler=%.3g", + len(self.memory_scalers), + len(self.index_bins), + self.max_memory_scaler, + ) return self.max_memory_scaler def next_batch(self) -> tuple[T | None, list[int]]: @@ -992,7 +1003,15 @@ def _get_first_batch(self) -> T: self.max_memory_scaler = self.max_memory_scaler * self.max_memory_padding newer_states = self._get_next_states() all_states.extend(newer_states) + logger.debug( + "InFlightAutoBatcher: estimated max_memory_scaler=%.3g", + self.max_memory_scaler, + ) + logger.info( + "InFlightAutoBatcher: starting with %d system(s) in first batch", + len(all_states), + ) return ts.concatenate_states(all_states) def next_batch( # noqa: C901 diff --git a/torch_sim/optimizers/fire.py b/torch_sim/optimizers/fire.py index d179a0e61..df1e22326 100644 --- a/torch_sim/optimizers/fire.py +++ b/torch_sim/optimizers/fire.py @@ -76,11 +76,13 @@ def fire_init( # Setup initial parameters dt_start_t = torch.as_tensor(dt_start, device=device, dtype=dtype) if dt_start_t.ndim == 0: - dt_start_t = dt_start_t.expand(n_systems) + # NOTE: clone needed as this is overwritten/assigned later by a masked_fill + dt_start_t = dt_start_t.expand(n_systems).clone() alpha_start_t = torch.as_tensor(alpha_start, device=device, dtype=dtype) if alpha_start_t.ndim == 0: - alpha_start_t = alpha_start_t.expand(n_systems) + # NOTE: clone needed as this is overwritten/assigned later by a masked_fill + alpha_start_t = alpha_start_t.expand(n_systems).clone() # FIRE-specific additional attributes fire_attrs = { diff --git a/torch_sim/runners.py b/torch_sim/runners.py index 08c205503..c8bafa6b7 100644 --- a/torch_sim/runners.py +++ b/torch_sim/runners.py @@ -293,6 +293,12 @@ def integrate[T: SimState]( # noqa: C901 unit_system = UnitSystem.metal initial_state: SimState = ts.initialize_state(system, model.device, model.dtype) + logger.info( + "integrate: n_systems=%d, n_steps=%d, integrator=%s", + initial_state.n_systems, + n_steps, + integrator, + ) dtype, device = initial_state.dtype, initial_state.device kTs = _normalize_temperature_tensor(temperature, n_steps, initial_state) kTs = kTs * unit_system.temperature @@ -388,8 +394,11 @@ def integrate[T: SimState]( # noqa: C901 if isinstance(batch_iterator, BinningAutoBatcher): reordered = batch_iterator.restore_original_order(final_states) # ty: ignore[invalid-argument-type] - return ts.concatenate_states(reordered) # ty: ignore[invalid-argument-type] + result = ts.concatenate_states(reordered) # ty: ignore[invalid-argument-type] + logger.info("integrate: complete, %d systems returned", result.n_systems) + return result + logger.info("integrate: complete") return state @@ -591,6 +600,12 @@ def optimize[T: OptimState]( # noqa: C901, PLR0915 convergence_fn = generate_energy_convergence_fn(energy_tol=1e-3) initial_state = ts.initialize_state(system, model.device, model.dtype) + logger.info( + "optimize: n_systems=%d, max_steps=%d, optimizer=%s", + initial_state.n_systems, + max_steps, + optimizer, + ) if isinstance(optimizer, Optimizer): init_fn, step_fn = OPTIM_REGISTRY[optimizer] elif ( @@ -683,6 +698,10 @@ def optimize[T: OptimState]( # noqa: C901, PLR0915 f"All systems have reached the maximum number of steps: {max_steps}.", stacklevel=2, ) + logger.warning( + "optimize: all systems reached max_steps=%d without converging", + max_steps, + ) break convergence_tensor = convergence_fn(state, last_energy) # ty: ignore[invalid-argument-type] @@ -699,7 +718,9 @@ def optimize[T: OptimState]( # noqa: C901, PLR0915 if autobatcher: final_states_ordered = autobatcher.restore_original_order(all_converged_states) - return ts.concatenate_states(final_states_ordered) + result = ts.concatenate_states(final_states_ordered) + logger.info("optimize: complete, %d systems returned", result.n_systems) + return result # Unreachable: _configure_in_flight_autobatcher always returns InFlightAutoBatcher raise RuntimeError( @@ -744,6 +765,7 @@ def static( list[dict[str, torch.Tensor]]: Maps of property names to tensors for all batches """ state: SimState = ts.initialize_state(system, model.device, model.dtype) + logger.info("static: n_systems=%d", state.n_systems) batch_iterator = _configure_batches_iterator(state, model, autobatcher=autobatcher) properties = ["potential_energy"] diff --git a/torch_sim/telemetry.py b/torch_sim/telemetry.py new file mode 100644 index 000000000..fd9e9d95b --- /dev/null +++ b/torch_sim/telemetry.py @@ -0,0 +1,98 @@ +"""Minimal logging for torch-sim. + +By default all log output is written to ``~/.torchsim/torch_sim.log`` (rotating, +up to 10 MB x 5 backups, JSON-lines format) **and** to the console at INFO level. +The file always captures DEBUG; the console level can be changed. + +The default log file accumulates entries across runs, which makes it easy to +review history, but means you need timestamps or ``--grep`` to isolate a single +run. To direct logs for a specific script to its own file, call +``configure_logging`` at the top of that script:: + + from torch_sim.telemetry import configure_logging, get_logger + + configure_logging(log_level="DEBUG", log_file="my_run.log") + log = get_logger(__name__) + log.info("starting simulation") + +``configure_logging`` is idempotent — calling it again replaces the handlers, so +you can safely reconfigure mid-script if needed. If you never call it explicitly, +logging auto-configures on the first ``get_logger()`` call using the defaults. +""" + +import json +import logging +import sys +from logging.handlers import RotatingFileHandler +from pathlib import Path + +from torch_sim import TORCH_SIM_CONFIG_DIR + + +_DEFAULT_LOG_FILE = TORCH_SIM_CONFIG_DIR / "torch_sim.log" +_configured = False + + +class _JSONFormatter(logging.Formatter): + """Emit one JSON object per line.""" + + def format(self, record: logging.LogRecord) -> str: + payload = { + "time": self.formatTime(record, "%Y-%m-%d %H:%M:%S"), + "level": record.levelname, + "name": record.name, + "message": record.getMessage(), + } + if record.exc_info: + payload["exception"] = self.formatException(record.exc_info) + return json.dumps(payload) + + +def configure_logging( + log_level: str = "INFO", + log_file: Path | str | None = None, +) -> None: + """Configure root logger with a console handler and a rotating JSON file handler. + + Args: + log_level: Minimum level shown on the console (``"DEBUG"``, ``"INFO"``, + ``"WARNING"``, ``"ERROR"``). The file always captures ``DEBUG`` regardless. + log_file: Path to the log file. Defaults to ``~/.torchsim/torch_sim.log``. + Pass a relative or absolute path to redirect a specific run. + + Notes: + - Safe to call multiple times; subsequent calls replace the existing handlers. + """ + global _configured # noqa: PLW0603 + log_file = Path(log_file) if log_file is not None else _DEFAULT_LOG_FILE + + root = logging.getLogger() + root.handlers.clear() + root.setLevel(logging.DEBUG) # handlers filter individually + + # Rotating JSON file — always DEBUG so nothing is lost + fh = RotatingFileHandler( + log_file, maxBytes=10 * 1024 * 1024, backupCount=5, encoding="utf-8" + ) + fh.setLevel(logging.DEBUG) + fh.setFormatter(_JSONFormatter()) + + # Human-readable console at the requested level + ch = logging.StreamHandler(sys.stdout) + ch.setLevel(log_level) + ch.setFormatter( + logging.Formatter( + "%(asctime)s %(levelname)-8s %(name)s: %(message)s", datefmt="%H:%M:%S" + ) + ) + + root.addHandler(fh) + root.addHandler(ch) + _configured = True + + +def get_logger(name: str | None = None) -> logging.Logger: + """Return a logger, auto-configuring on first call.""" + if not _configured: + configure_logging() + return logging.getLogger(name) From 9cffdcd1c9511b33716a1757e2abd02441ccc910 Mon Sep 17 00:00:00 2001 From: Rhys Goodall Date: Mon, 2 Mar 2026 11:48:59 -0500 Subject: [PATCH 2/4] Update torch_sim/runners.py Co-authored-by: Orion Cohen <27712051+orionarcher@users.noreply.github.com> Signed-off-by: Rhys Goodall --- torch_sim/runners.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torch_sim/runners.py b/torch_sim/runners.py index c8bafa6b7..5ad9913d6 100644 --- a/torch_sim/runners.py +++ b/torch_sim/runners.py @@ -699,7 +699,7 @@ def optimize[T: OptimState]( # noqa: C901, PLR0915 stacklevel=2, ) logger.warning( - "optimize: all systems reached max_steps=%d without converging", + "optimize: all systems in batch reached max_steps=%d without converging", max_steps, ) break From aa9501985788e67450edaf08a8eb0028258a2081 Mon Sep 17 00:00:00 2001 From: Rhys Goodall Date: Mon, 2 Mar 2026 13:01:11 -0500 Subject: [PATCH 3/4] lint --- torch_sim/runners.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/torch_sim/runners.py b/torch_sim/runners.py index 5ad9913d6..00afceff2 100644 --- a/torch_sim/runners.py +++ b/torch_sim/runners.py @@ -197,11 +197,13 @@ def _normalize_temperature_tensor( return temps.expand(n_steps) if initial_state.n_systems == n_steps: - warnings.warn( + msg = ( "n_systems is equal to n_steps. Interpreting temperature array of length " - "n_systems as temperatures for each system, broadcasted over steps.", - stacklevel=2, + "n_systems as temperatures for each system, broadcasted over steps." ) + warnings.warn(msg, stacklevel=2) + logger.warning(msg) + return temps.expand(n_steps) if temps.shape[0] == initial_state.n_systems: if temps.ndim == 2: @@ -694,14 +696,11 @@ def optimize[T: OptimState]( # noqa: C901, PLR0915 step[autobatcher.current_idx] += 1 exceeded_max_steps = step > max_steps if exceeded_max_steps.all(): - warnings.warn( - f"All systems have reached the maximum number of steps: {max_steps}.", - stacklevel=2, - ) - logger.warning( - "optimize: all systems in batch reached max_steps=%d without converging", - max_steps, + msg = ( + f"All systems have reached the maximum number of steps: {max_steps}." ) + warnings.warn(msg, stacklevel=2) + logger.warning(msg) break convergence_tensor = convergence_fn(state, last_energy) # ty: ignore[invalid-argument-type] From e877954b492fc8b68e749420008dda0164e9aaf2 Mon Sep 17 00:00:00 2001 From: Rhys Goodall Date: Mon, 2 Mar 2026 13:06:43 -0500 Subject: [PATCH 4/4] maint: log all relevant warnings --- torch_sim/constraints.py | 19 +++++++++++-------- torch_sim/integrators/md.py | 12 +++++++----- torch_sim/integrators/npt.py | 20 ++++++++++++-------- torch_sim/trajectory.py | 10 +++++++--- 4 files changed, 37 insertions(+), 24 deletions(-) diff --git a/torch_sim/constraints.py b/torch_sim/constraints.py index d1557a4b4..e33d1e0b0 100644 --- a/torch_sim/constraints.py +++ b/torch_sim/constraints.py @@ -10,6 +10,7 @@ from __future__ import annotations +import logging import math import warnings from abc import ABC, abstractmethod @@ -18,6 +19,8 @@ import torch +logger = logging.getLogger(__name__) + if TYPE_CHECKING: from torch_sim.state import SimState @@ -694,22 +697,22 @@ def validate_constraints(constraints: list[Constraint], state: SimState) -> None all_indices = torch.cat([c.atom_idx for c in indexed_constraints]) unique_indices = torch.unique(all_indices) if len(unique_indices) < len(all_indices): - warnings.warn( + msg = ( "Multiple constraints are acting on the same atoms. " - "This may lead to unexpected behavior.", - UserWarning, - stacklevel=3, + "This may lead to unexpected behavior." ) + warnings.warn(msg, UserWarning, stacklevel=3) + logger.warning(msg) # Warn about COM constraint with fixed atoms if has_com_constraint and indexed_constraints: - warnings.warn( + msg = ( "Using FixCom together with other constraints may lead to " "unexpected behavior. The center of mass constraint is applied " - "to all atoms, including those that may be constrained by other means.", - UserWarning, - stacklevel=3, + "to all atoms, including those that may be constrained by other means." ) + warnings.warn(msg, UserWarning, stacklevel=3) + logger.warning(msg) class FixSymmetry(SystemConstraint): diff --git a/torch_sim/integrators/md.py b/torch_sim/integrators/md.py index 47aafcf40..5a8c87743 100644 --- a/torch_sim/integrators/md.py +++ b/torch_sim/integrators/md.py @@ -1,5 +1,6 @@ """Core molecular dynamics state and operations.""" +import logging import warnings from collections.abc import Callable from dataclasses import dataclass @@ -13,6 +14,9 @@ from torch_sim.units import MetalUnits +logger = logging.getLogger(__name__) + + @dataclass(kw_only=True) class MDState(SimState): """State information for molecular dynamics simulations. @@ -236,11 +240,9 @@ def velocity_verlet[T: MDState]( state: T, dt: float | torch.Tensor, model: ModelInterface ) -> T: """Deprecated alias for velocity_verlet_step.""" - warnings.warn( - "velocity_verlet is deprecated. Use velocity_verlet_step instead.", - DeprecationWarning, - stacklevel=2, - ) + msg = "velocity_verlet is deprecated. Use velocity_verlet_step instead." + warnings.warn(msg, DeprecationWarning, stacklevel=2) + logger.warning(msg) return velocity_verlet_step(state=state, dt=dt, model=model) diff --git a/torch_sim/integrators/npt.py b/torch_sim/integrators/npt.py index eb6cb4677..80df895bf 100644 --- a/torch_sim/integrators/npt.py +++ b/torch_sim/integrators/npt.py @@ -1,5 +1,6 @@ """Implementations of NPT integrators.""" +import logging import warnings from collections.abc import Callable from dataclasses import dataclass @@ -23,6 +24,9 @@ from torch_sim.typing import StateDict +logger = logging.getLogger(__name__) + + def _randn_for_state(state: MDState, shape: torch.Size | tuple[int, ...]) -> torch.Tensor: """Sample standard normal noise on the state's device/dtype using state RNG.""" return torch.randn(shape, device=state.device, dtype=state.dtype, generator=state.rng) @@ -612,13 +616,13 @@ def npt_langevin_init( if state.constraints: # warn if constraints are present - warnings.warn( + msg = ( "Constraints are present in the system. " "Make sure they are compatible with NPT Langevin dynamics." - "We recommend not using constraints with NPT dynamics for now.", - UserWarning, - stacklevel=3, + "We recommend not using constraints with NPT dynamics for now." ) + warnings.warn(msg, UserWarning, stacklevel=3) + logger.warning(msg) # Create the initial state return NPTLangevinState.from_state( @@ -1427,13 +1431,13 @@ def npt_nose_hoover_init( if state.constraints: # warn if constraints are present - warnings.warn( + msg = ( "Constraints are present in the system. " "Make sure they are compatible with NPT Nosé Hoover dynamics." - "We recommend not using constraints with NPT dynamics for now.", - UserWarning, - stacklevel=3, + "We recommend not using constraints with NPT dynamics for now." ) + warnings.warn(msg, UserWarning, stacklevel=3) + logger.warning(msg) # Create initial state return NPTNoseHooverState.from_state( diff --git a/torch_sim/trajectory.py b/torch_sim/trajectory.py index 27d1785fa..290351496 100644 --- a/torch_sim/trajectory.py +++ b/torch_sim/trajectory.py @@ -29,6 +29,7 @@ import copy import inspect +import logging import pathlib import warnings from collections.abc import Callable, Mapping, Sequence @@ -43,6 +44,8 @@ from torch_sim.state import SimState +logger = logging.getLogger(__name__) + if TYPE_CHECKING: from ase import Atoms from ase.io.trajectory import TrajectoryReader @@ -547,11 +550,12 @@ def __init__( self.get_steps(name)[-1] > self.last_step for name in self.array_registry ) if inconsistent_step: - warnings.warn( + msg = ( "Inconsistent last steps detected in trajectory arrays. " - "Truncating all arrays to the `positions` array's last step.", - stacklevel=2, + "Truncating all arrays to the `positions` array's last step." ) + warnings.warn(msg, UserWarning, stacklevel=2) + logger.warning(msg) self.truncate_to_step(self.last_step) def _initialize_header(self, metadata: dict[str, str] | None = None) -> None: