First version of the code which can produce plots! Look totally wrong: some major issue in the physics.

Author: jarvist

examples/pyrazine/gsb_and_se_gnuplot.jl

@@ -0,0 +1,47 @@
+using Janus
+using Gnuplot
+
+# Load trajectory data (same loader for GSB and SE!)
+traj_gs = load_trajectory("gsb/gsb_ene.dat", "gsb/gsb_osc.dat")
+
+# Laser pulse (5 fs Gaussian, no dispersion)
+laser = LaserPulse(τ=5.0)
+
+# Experimental parameters
+ω_pump  = 3.0                             # eV
+ω_probe = collect(range(1, 6; length=200)) # eV grid
+
+# GSB spectrum
+S_gsb = spectrum(GSB(), traj_gs, laser, ω_pump, ω_probe)
+
+@show S_gsb
+
+@gp "set xlabel 'Probe energy (eV)'" "set ylabel 'Signal (arb.)'" :-
+@gp :- "set xrange [0:50]"
+@gp :- "set yrange [1:200]"
+@gp :- S_gsb' "with image title 'GSB'"
+Gnuplot.save("gsb.png", term="pngcairo size 550,350 fontscale 0.8")
+
+# SE -- same loader, different signal type
+traj_se = load_trajectory("se/se_ene.dat", "se/se_osc.dat")
+S_se = spectrum(SE(), traj_se, laser, ω_pump, ω_probe)
+
+@show S_se
+
+# Time-resolved 2D map
+@gp "set xlabel 'Time step'" "set ylabel 'Probe index'" :-
+@gp :- "set xrange [0:50]"
+@gp :- "set yrange [1:200]"
+@gp :- S_se' "with image title 'SE'"
+Gnuplot.save("se.png", term="pngcairo size 550,350 fontscale 0.8")
+
+# different pump and probe lasers
+pump_laser  = LaserPulse(τ=5.0, envelope=Gaussian())
+probe_laser = LaserPulse(τ=10.0, envelope=Sech())
+S_nd = spectrum(GSB(), traj_gs, pump_laser, ω_pump, probe_laser, ω_probe)
+
+@gp "set xlabel 'Time step'" "set ylabel 'Probe index'" :-
+@gp :- "set xrange [0:50]"
+@gp :- "set yrange [1:200]"
+@gp :- S_nd' "with image title 'different pump and probe lasers'"
+Gnuplot.save("gsb_weird_laser.png", term="pngcairo size 550,350 fontscale 0.8")

src/Janus.jl

@@ -1,5 +1,27 @@
 module Janus
 
-# Write your package code here.
+# Janus.jl
+# Simple implementation of Doorway-Window approach for time-dependent spectrscopy 
+
+using DelimitedFiles
+using LinearAlgebra: mul!
+
+
+include("constants.jl")
+export eV2Ha, Ha2eV, fs2au, au2fs, nm2eV, eV2nm, eV2invcm, invcm2eV
+
+include("types.jl")
+export SignalType, GSB, SE, ESA
+export Envelope, Gaussian, Lorentzian, Sech
+export Trajectory, LaserPulse, Dephasing
+
+include("io.jl")
+export load_trajectory
+export delta_energy, n_states, n_steps
+
+include("dw.jl")
+export dipole_moment, spectral_envelope
+export doorway, window, spectrum
 
 end
+

src/constants.jl

@@ -0,0 +1,20 @@
+const Ha2eV  = 27.2113862459      # Hartree -> eV
+const eV2Ha  = 1.0 / Ha2eV        # eV -> Hartree
+const fs2au  = 41.341374575751    # fs -> a.u. time
+const au2fs  = 1.0 / fs2au        # a.u. time -> fs
+
+# Careful! Energy is reciprocal of wavelength; these are dangerously named
+const eV2nm = 1239.84
+const nm2eV = 1.0 / eV2nm
+
+const eV2invcm = 8065.5439373492107
+const invcm2eV = 1.0 / eV2invcm
+
+# Factorial constants for Taylor expansions; a little bit OTT perhaps
+# I wonder if this actually gets any speed increase? Only used in Gaussian wavepackets? Are
+# those called repeatedly for the Window?
+const _inv_fact2 = 1 / factorial(2)
+const _inv_fact3 = 1 / factorial(3)
+const _inv_fact4 = 1 / factorial(4)
+const _inv_fact5 = 1 / factorial(5)
+

src/dw.jl

@@ -0,0 +1,103 @@
+# magic broadcast with @. 
+# factor 3//2 due to isotropic overage of dipole
+dipole_moment_sq(f, ΔE) = @. 3//2 * f / ΔE
+
+# DOORWAYS
+
+# "Beware of Doors." - Neverwhere
+
+# GSB 
+function doorway(::GSB, traj::Trajectory, pump::LaserPulse, ω_pump::Real, frame::Int)
+    ΔE  = delta_energy(traj, frame)
+    μ²  = dipole_moment_sq(traj.oscillators[:, frame], ΔE)
+    Δω  = @. ω_pump - ΔE
+    τ_au = pump.τ * fs2au
+    E   = spectral_envelope(Δω, τ_au, pump.envelope, pump)
+    return sum(@. E^2 * μ²)   # returns a scalar, no mutation
+end
+
+# SE
+function doorway(::SE, traj::Trajectory, pump::LaserPulse, ω_pump::Real, frame::Int)
+    s   = traj.active_state[frame]
+    ΔE  = traj.energies[s + 1, frame] - traj.energies[1, frame]
+    μ²  = dipole_moment_sq(traj.oscillators[s, frame], ΔE)
+    Δω  = ω_pump - ΔE
+    τ_au = pump.τ * fs2au
+    E   = spectral_envelope(Δω, τ_au, pump.envelope, pump)
+    return E^2 * μ²   # scalar, single active state
+end
+
+# WINDOWS
+
+# GSB
+function window(::GSB, traj::Trajectory, probe::LaserPulse, ω_probe::AbstractVector)
+    ΔE  = delta_energy(traj)                              # (n_states-1) x n_steps
+    μ²  = dipole_moment_sq(traj.oscillators, ΔE)          # (n_states-1) x n_steps
+    n_t = n_steps(traj)
+    n_ω = length(ω_probe)
+    τ_au = probe.τ * fs2au
+
+    W = zeros(n_t, n_ω)
+    for (j, ω) in enumerate(ω_probe)
+        Δω = @. ω .- ΔE                                   # (n_states-1) x n_steps
+        E  = spectral_envelope(Δω, τ_au, probe.envelope, probe)
+        @views W[:, j] .= vec(sum(@. E^2 * μ²; dims=1))  # sum over states
+    end
+    return W
+end
+
+# SE
+function window(::SE, traj::Trajectory, probe::LaserPulse, ω_probe::AbstractVector)
+    n_t = n_steps(traj)
+    n_ω = length(ω_probe)
+    τ_au = probe.τ * fs2au
+
+    # Extract active-state ΔE and oscillator for each frame
+    ΔE_active = Vector{Float64}(undef, n_t)
+    f_active  = Vector{Float64}(undef, n_t)
+    for t in 1:n_t
+        s = traj.active_state[t]
+        ΔE_active[t] = traj.energies[s + 1, t] - traj.energies[1, t]
+        f_active[t]  = traj.oscillators[s, t]
+    end
+
+    μ² = dipole_moment_sq(f_active, ΔE_active)
+    # Guard against Inf/NaN from zero ΔE (trajectory at ground state)
+    μ² .= ifelse.(isfinite.(μ²), μ², 0.0)
+
+    W = zeros(n_t, n_ω)
+    for (j, ω) in enumerate(ω_probe)
+        Δω = @. ω - ΔE_active
+        E  = spectral_envelope(Δω, τ_au, probe.envelope, probe)
+        @views W[:, j] .= @. E^2 * μ²
+    end
+    return W
+end
+
+# OK, generate the spectrum!
+
+# Convenience function for same LaserPulse in both pump and probe
+function spectrum(sig::SignalType, traj::Trajectory,
+                  laser::LaserPulse, ω_pump::Real, ω_probe::AbstractVector)
+    spectrum(sig, traj, laser, ω_pump, laser, ω_probe)
+end
+
+"""
+Pump-probe spectrum: doorway (n_steps,) .* window (n_steps, n_probe).
+
+User-facing frequencies ω_pump, ω_probe are in **eV**; internally converted
+to Hartree to match trajectory energies.
+
+Returns: (n_steps, n_probe) matrix. Sum over time steps for integrated signal.
+"""
+function spectrum(sig::SignalType, traj::Trajectory,
+                  pump::LaserPulse, ω_pump::Real,
+                  probe::LaserPulse, ω_probe::AbstractVector)
+    ω_pu_au = ω_pump * eV2Ha
+    ω_pr_au = ω_probe .* eV2Ha
+    d = [doorway(sig, traj, pump, ω_pu_au, t) for t in 1:n_steps(traj)]
+    W = window(sig, traj, probe, ω_pr_au)
+    return d .* W
+end
+
+

src/io.jl

@@ -0,0 +1,40 @@
+# ZagHop ASCII format (?) currently we just have the examples with Wavemixings to test with
+
+function load_trajectory(energy_file::AbstractString, oscillator_file::AbstractString)
+    # Energy file: #header, then columns [time, cst, etot, epot, qe001, qe002, ...]
+    raw = readdlm(energy_file; comments=true, comment_char='#')
+
+    time         = @view raw[:, 1]
+    active_state = Int.(@view raw[:, 2])
+
+    # Columns 5:end are the adiabatic state energies (skip time, cst, etot, epot)
+    # Transpose so rows = states, columns = time steps
+    energies = permutedims(@view raw[:, 5:end])
+
+    # Handle duplicate time steps (surface hops produce repeated entries)
+    keep = _unique_timesteps(time, active_state)
+
+    # Oscillator file: no header, columns are oscillator strengths per transition
+    osc_raw = readdlm(oscillator_file; comments=true, comment_char='#')
+    oscillators = permutedims(osc_raw)
+
+    Trajectory(
+        time[keep],
+        active_state[keep],
+        energies[:, keep],
+        max.(oscillators[:, keep], 0.0),  # clamp negative oscillator strengths
+    )
+end
+
+function _unique_timesteps(time, active_state)
+    # When a surface hop occurs, ZagHop seems to write the same time step twice (old state
+    # then new state). Keep only the last entry for each time.
+    keep = trues(length(time))
+    for i in 1:(length(time) - 1)
+        if time[i] == time[i + 1]
+            keep[i] = false
+        end
+    end
+    return keep
+end
+

src/types.jl

@@ -0,0 +1,95 @@
+# Abstract types for dispatch
+
+abstract type SignalType end
+struct GSB <: SignalType end   # Ground State Bleach
+struct SE  <: SignalType end   # Stimulated Emission
+struct ESA <: SignalType end   # Excited State Absorption
+
+abstract type Envelope end
+struct Gaussian   <: Envelope end
+struct Lorentzian <: Envelope end
+struct Sech       <: Envelope end
+
+
+
+"""
+Trajectory data from a non-adiabatic (surface-hopping kind presumed) simulation.
+
+Fields:
+- `time::Vector{Float64}`: simulation times (fs)
+- `active_state::Vector{Int}`: occupied electronic state (density?) at each step
+- `energies::Matrix{Float64}`: adiabatic energies (n_states x n_steps), Hartree
+- `oscillators::Matrix{Float64}`: oscillator strengths ((n_states-1) x n_steps)
+"""
+struct Trajectory
+    time::Vector{Float64}
+    active_state::Vector{Int}
+    energies::Matrix{Float64}
+    oscillators::Matrix{Float64}
+end
+
+n_states(t::Trajectory) = size(t.energies, 1)
+n_steps(t::Trajectory)  = size(t.energies, 2)
+
+"""ΔE_e = E_e - E_ground for all excited states at one frame."""
+delta_energy(traj::Trajectory, frame::Int) =
+    @views traj.energies[2:end, frame] .- traj.energies[1, frame]
+
+"""ΔE matrix for all frames: (n_states-1) x n_steps."""
+delta_energy(traj::Trajectory) =
+    @views traj.energies[2:end, :] .- traj.energies[1:1, :]
+
+
+"""
+Laser pulse characteristics (envelope shape, duration, dispersion).
+(Nb: central frequency passed elsewhere; this is just arb pulse!)
+
+Fields:
+- `envelope::E`: spectral envelope shape (Gaussian, Lorentzian, Sech, etc.)
+- `τ::Float64`: pulse duration (fs)
+- `D2`--`D5`: dispersion coefficients (fs^n), default 0
+"""
+@kwdef struct LaserPulse{E<:Envelope}
+    envelope::E = Gaussian()
+    τ::Float64
+    D2::Float64 = 0.0
+    D3::Float64 = 0.0
+    D4::Float64 = 0.0
+    D5::Float64 = 0.0
+end
+
+"""Dephasing rate for dispersed and 2D calculations."""
+@kwdef struct Dephasing
+    γ::Float64 = 0.01  # eV
+end
+# not used yet
+
+# Nb: spectra envelope stuff needs to be down here, as depends on LaserPulse having been ready by Julia
+"""
+    spectral_envelope(Δω, τ_au, envelope, pulse)
+
+Spectral field envelope E(Δω, τ) in the frequency domain.
+
+Returns the field amplitude. Signal formulas use E² (proportional to intensity).
+All envelopes include τ_au prefactor for normalisation. 
+"""
+function spectral_envelope(Δω, τ_au::Real, ::Gaussian, p::LaserPulse)
+    @. τ_au * exp(
+        -(Δω * τ_au)^2 / 4 - (
+            _inv_fact2 * p.D2 * Δω^2 +
+            _inv_fact3 * p.D3 * Δω^3 +
+            _inv_fact4 * p.D4 * Δω^4 +
+            _inv_fact5 * p.D5 * Δω^5
+        )
+    )
+end
+
+# Lorentzian
+function spectral_envelope(Δω, τ_au::Real, ::Lorentzian, ::LaserPulse)
+    @. τ_au * exp(-2 * abs(Δω) * τ_au)
+end
+
+# Sech
+function spectral_envelope(Δω, τ_au::Real, ::Sech, ::LaserPulse)
+    @. τ_au * sech(π * Δω * τ_au / 2)
+end