fennol.md.thermostats

  1import numpy as np
  2import flax.linen as nn
  3import jax
  4import jax.numpy as jnp
  5import math
  6import optax
  7import os
  8import pickle
  9
 10from .utils import us 
 11from ..utils import Counter,read_tinker_interval
 12from ..utils.deconvolution import (
 13    deconvolute_spectrum,
 14    kernel_lorentz_pot,
 15    kernel_lorentz,
 16)
 17from ..utils.periodic_table import PERIODIC_TABLE
 18
 19
 20def get_thermostat(simulation_parameters, dt, system_data, fprec, rng_key=None, restart_data={}):
 21    state = {}
 22    postprocess = None
 23    
 24
 25    thermostat_name = str(simulation_parameters.get("thermostat", "LGV")).upper()
 26    """@keyword[fennol_md] thermostat
 27    Thermostat type. Options: 'NVE', 'LGV', 'NOSE', 'ADQTB'.
 28    Default: "LGV"
 29    """
 30    compute_thermostat_energy = simulation_parameters.get(
 31        "include_thermostat_energy", False
 32    )
 33    """@keyword[fennol_md] include_thermostat_energy
 34    Include thermostat energy in total energy calculations.
 35    Default: False
 36    """
 37
 38    kT = system_data.get("kT", None)
 39    nbeads = system_data.get("nbeads", None)
 40    mass = system_data["mass"]
 41    gamma = simulation_parameters.get("gamma", 1.0 / us.THZ)
 42    """@keyword[fennol_md] gamma
 43    Friction coefficient for Langevin thermostat.
 44    Default: 1.0 ps^-1
 45    """
 46    species = system_data["species"]
 47
 48    if nbeads is not None:
 49        trpmd_lambda = simulation_parameters.get("trpmd_lambda", 1.0)
 50        """@keyword[fennol_md] trpmd_lambda
 51        Lambda parameter for TRPMD (Thermostatted Ring Polymer MD).
 52        Default: 1.0
 53        """
 54        gamma = np.maximum(trpmd_lambda * system_data["omk"], gamma)
 55
 56    if thermostat_name in ["LGV", "LANGEVIN", "FFLGV"]:
 57        assert rng_key is not None, "rng_key must be provided for QTB thermostat"
 58        assert kT is not None, "kT must be specified for QTB thermostat"
 59        assert gamma is not None, "gamma must be specified for QTB thermostat"
 60        rng_key, v_key = jax.random.split(rng_key)
 61        if nbeads is None:
 62            a1 = math.exp(-gamma * dt)
 63            a2 = jnp.asarray(((1 - a1 * a1) * kT / mass[:, None]) ** 0.5, dtype=fprec)
 64            vel = (
 65                jax.random.normal(v_key, (mass.shape[0], 3), dtype=fprec)
 66                * (kT / mass[:, None]) ** 0.5
 67            )
 68        else:
 69            if isinstance(gamma, float):
 70                gamma = np.array([gamma] * nbeads)
 71            assert isinstance(
 72                gamma, np.ndarray
 73            ), "gamma must be a float or a numpy array"
 74            assert gamma.shape[0] == nbeads, "gamma must have the same length as nbeads"
 75            a1 = np.exp(-gamma * dt)[:, None, None]
 76            a2 = jnp.asarray(
 77                ((1 - a1 * a1) * kT / mass[None, :, None]) ** 0.5, dtype=fprec
 78            )
 79            vel = (
 80                jax.random.normal(v_key, (nbeads, mass.shape[0], 3), dtype=fprec)
 81                * (kT / mass[:, None]) ** 0.5
 82            )
 83
 84        state["rng_key"] = rng_key
 85        if compute_thermostat_energy:
 86            state["thermostat_energy"] = 0.0
 87        if thermostat_name == "FFLGV":
 88            def thermostat(vel, state):
 89                rng_key, noise_key = jax.random.split(state["rng_key"])
 90                noise = jax.random.normal(noise_key, vel.shape, dtype=vel.dtype)
 91                norm_vel = jnp.linalg.norm(vel, axis=-1, keepdims=True)
 92                dirvel = vel / norm_vel
 93                if compute_thermostat_energy:
 94                    v2 = (vel**2).sum(axis=-1)
 95                vel = a1 * vel + a2 * noise
 96                new_norm_vel = jnp.linalg.norm(vel, axis=-1, keepdims=True)
 97                vel = dirvel * new_norm_vel
 98                new_state = {**state, "rng_key": rng_key}
 99                if compute_thermostat_energy:
100                    v2new = (vel**2).sum(axis=-1)
101                    new_state["thermostat_energy"] = (
102                        state["thermostat_energy"] + 0.5 * (mass * (v2 - v2new)).sum()
103                    )
104
105                return vel, new_state
106
107        else:
108            def thermostat(vel, state):
109                rng_key, noise_key = jax.random.split(state["rng_key"])
110                noise = jax.random.normal(noise_key, vel.shape, dtype=vel.dtype)
111                if compute_thermostat_energy:
112                    v2 = (vel**2).sum(axis=-1)
113                vel = a1 * vel + a2 * noise
114                new_state = {**state, "rng_key": rng_key}
115                if compute_thermostat_energy:
116                    v2new = (vel**2).sum(axis=-1)
117                    new_state["thermostat_energy"] = (
118                        state["thermostat_energy"] + 0.5 * (mass * (v2 - v2new)).sum()
119                    )
120                return vel, new_state
121
122    elif thermostat_name in ["BUSSI"]:
123        assert rng_key is not None, "rng_key must be provided for QTB thermostat"
124        assert kT is not None, "kT must be specified for QTB thermostat"
125        assert gamma is not None, "gamma must be specified for QTB thermostat"
126        assert nbeads is None, "Bussi thermostat is not compatible with PIMD"
127        rng_key, v_key = jax.random.split(rng_key)
128
129        a1 = math.exp(-gamma * dt)
130        a2 = (1 - a1) * kT
131        vel = (
132            jax.random.normal(v_key, (mass.shape[0], 3), dtype=fprec)
133            * (kT / mass[:, None]) ** 0.5
134        )
135
136        state["rng_key"] = rng_key
137        if compute_thermostat_energy:
138            state["thermostat_energy"] = 0.0
139
140        def thermostat(vel, state):
141            rng_key, noise_key = jax.random.split(state["rng_key"])
142            new_state = {**state, "rng_key": rng_key}
143            noise = jax.random.normal(noise_key, vel.shape, dtype=vel.dtype)
144            R2 = jnp.sum(noise**2)
145            R1 = noise[0, 0]
146            c = a2 / (mass[:, None] * vel**2).sum()
147            d = (a1 * c) ** 0.5
148            scale = (a1 + c * R2 + 2 * d * R1) ** 0.5
149            if compute_thermostat_energy:
150                dek = 0.5 * (mass[:, None] * vel**2).sum() * (scale**2 - 1)
151                new_state["thermostat_energy"] = state["thermostat_energy"] + dek
152            return scale * vel, new_state
153
154    elif thermostat_name in [
155        "GD",
156        "DESCENT",
157        "GRADIENT_DESCENT",
158        "MIN",
159        "MINIMIZE",
160    ]:
161        assert nbeads is None, "Gradient descent is not compatible with PIMD"
162        a1 = math.exp(-gamma * dt)
163
164        if nbeads is None:
165            vel = jnp.zeros((mass.shape[0], 3), dtype=fprec)
166        else:
167            vel = jnp.zeros((nbeads, mass.shape[0], 3), dtype=fprec)
168
169        def thermostat(vel, state):
170            return a1 * vel, state
171
172    elif thermostat_name in ["NVE", "NONE"]:
173        if nbeads is None:
174            vel = (
175                jax.random.normal(rng_key, (mass.shape[0], 3), dtype=fprec)
176                * (kT / mass[:, None]) ** 0.5
177            )
178            kTsys = jnp.sum(mass[:, None] * vel**2) / (mass.shape[0] * 3)
179            vel = vel * (kT / kTsys) ** 0.5
180        else:
181            vel = (
182                jax.random.normal(rng_key, (nbeads, mass.shape[0], 3), dtype=fprec)
183                * (kT / mass[None, :, None]) ** 0.5
184            )
185            kTsys = jnp.sum(mass[None, :, None] * vel**2, axis=(1, 2)) / (
186                mass.shape[0] * 3
187            )
188            vel = vel * (kT / kTsys[:, None, None]) ** 0.5
189        thermostat = lambda x, s: (x, s)
190
191    elif thermostat_name in ["NOSE", "NOSEHOOVER", "NOSE_HOOVER"]:
192        assert gamma is not None, "gamma must be specified for QTB thermostat"
193        ndof = mass.shape[0] * 3
194        nkT = ndof * kT
195        nose_mass = nkT / gamma**2
196        assert nbeads is None, "Nose-Hoover is not compatible with PIMD"
197        state["nose_s"] = 0.0
198        state["nose_v"] = 0.0
199        if compute_thermostat_energy:
200            state["thermostat_energy"] = 0.0
201        print(
202            "# WARNING: Nose-Hoover thermostat is not well tested yet. Energy conservation is not guaranteed."
203        )
204        vel = (
205            jax.random.normal(rng_key, (mass.shape[0], 3), dtype=fprec)
206            * (kT / mass[:, None]) ** 0.5
207        )
208
209        def thermostat(vel, state):
210            nose_s = state["nose_s"]
211            nose_v = state["nose_v"]
212            kTsys = jnp.sum(mass[:, None] * vel**2)
213            nose_v = nose_v + (0.5 * dt / nose_mass) * (kTsys - nkT)
214            nose_s = nose_s + dt * nose_v
215            vel = jnp.exp(-nose_v * dt) * vel
216            kTsys = jnp.sum(mass[:, None] * vel**2)
217            nose_v = nose_v + (0.5 * dt / nose_mass) * (kTsys - nkT)
218            new_state = {**state, "nose_s": nose_s, "nose_v": nose_v}
219
220            if compute_thermostat_energy:
221                new_state["thermostat_energy"] = (
222                    nkT * nose_s + (0.5 * nose_mass) * nose_v**2
223                )
224            return vel, new_state
225
226    elif thermostat_name in ["QTB", "ADQTB"]:
227        assert nbeads is None, "QTB is not compatible with PIMD"
228        assert rng_key is not None, "rng_key must be provided for QTB thermostat"
229        assert kT is not None, "kT must be specified for QTB thermostat"
230        assert gamma is not None, "gamma must be specified for QTB thermostat"
231        assert species is not None, "species must be provided for QTB thermostat"
232        rng_key, v_key = jax.random.split(rng_key)
233        vel = (
234            jax.random.normal(v_key, (mass.shape[0], 3), dtype=fprec)
235            * (kT / mass[:, None]) ** 0.5
236        )
237
238        thermostat, postprocess, qtb_state = initialize_qtb(
239            simulation_parameters,
240            system_data,
241            fprec=fprec,
242            dt=dt,
243            mass=mass,
244            gamma=gamma,
245            kT=kT,
246            species=species,
247            rng_key=rng_key,
248            adaptive=thermostat_name.startswith("AD"),
249            compute_thermostat_energy=compute_thermostat_energy,
250        )
251        state = {**state, **qtb_state}
252
253    elif thermostat_name in ["ANNEAL", "ANNEALING"]:
254        assert rng_key is not None, "rng_key must be provided for QTB thermostat"
255        assert kT is not None, "kT must be specified for QTB thermostat"
256        assert gamma is not None, "gamma must be specified for QTB thermostat"
257        assert nbeads is None, "ANNEAL is not compatible with PIMD"
258        a1 = math.exp(-gamma * dt)
259        a2 = jnp.asarray(((1 - a1 * a1) * kT / mass[:, None]) ** 0.5, dtype=fprec)
260
261        anneal_parameters = simulation_parameters.get("annealing", {})
262        """@keyword[fennol_md] annealing
263        Parameters for simulated annealing schedule configuration.
264        Required for ANNEAL/ANNEALING thermostat
265        """
266        init_factor = anneal_parameters.get("init_factor", 1.0 / 25.0)
267        """@keyword[fennol_md] annealing/init_factor
268        Initial temperature factor for annealing schedule.
269        Default: 0.04 (1/25)
270        """
271        assert init_factor > 0.0, "init_factor must be positive"
272        final_factor = anneal_parameters.get("final_factor", 1.0 / 10000.0)
273        """@keyword[fennol_md] annealing/final_factor
274        Final temperature factor for annealing schedule.
275        Default: 0.0001 (1/10000)
276        """
277        assert final_factor > 0.0, "final_factor must be positive"
278        nsteps = simulation_parameters.get("nsteps")
279        """@keyword[fennol_md] nsteps
280        Total number of simulation steps for annealing schedule calculation.
281        Required parameter
282        """
283        anneal_steps = anneal_parameters.get("anneal_steps", 1.0)
284        """@keyword[fennol_md] annealing/anneal_steps
285        Fraction of total steps for annealing process.
286        Default: 1.0
287        """
288        assert (
289            anneal_steps < 1.0 and anneal_steps > 0.0
290        ), "warmup_steps must be between 0 and nsteps"
291        pct_start = anneal_parameters.get("warmup_steps", 0.3)
292        """@keyword[fennol_md] annealing/warmup_steps
293        Fraction of annealing steps for warmup phase.
294        Default: 0.3
295        """
296        assert (
297            pct_start < 1.0 and pct_start > 0.0
298        ), "warmup_steps must be between 0 and nsteps"
299
300        anneal_type = anneal_parameters.get("type", "cosine").lower()
301        """@keyword[fennol_md] annealing/type
302        Type of annealing schedule (linear, cosine_onecycle).
303        Default: cosine
304        """
305        if anneal_type == "linear":
306            schedule = optax.linear_onecycle_schedule(
307                peak_value=1.0,
308                div_factor=1.0 / init_factor,
309                final_div_factor=1.0 / final_factor,
310                transition_steps=int(anneal_steps * nsteps),
311                pct_start=pct_start,
312                pct_final=1.0,
313            )
314        elif anneal_type == "cosine_onecycle":
315            schedule = optax.cosine_onecycle_schedule(
316                peak_value=1.0,
317                div_factor=1.0 / init_factor,
318                final_div_factor=1.0 / final_factor,
319                transition_steps=int(anneal_steps * nsteps),
320                pct_start=pct_start,
321            )
322        else:
323            raise ValueError(f"Unknown anneal_type {anneal_type}")
324
325        state["rng_key"] = rng_key
326        state["istep_anneal"] = 0
327
328        rng_key, v_key = jax.random.split(rng_key)
329        Tscale = schedule(0)
330        print(f"# ANNEAL: initial temperature = {Tscale*kT*us.KELVIN:.3e} K")
331        vel = (
332            jax.random.normal(v_key, (mass.shape[0], 3), dtype=fprec)
333            * (kT * Tscale / mass[:, None]) ** 0.5
334        )
335
336        def thermostat(vel, state):
337            rng_key, noise_key = jax.random.split(state["rng_key"])
338            noise = jax.random.normal(noise_key, vel.shape, dtype=vel.dtype)
339
340            Tscale = schedule(state["istep_anneal"]) ** 0.5
341            vel = a1 * vel + a2 * Tscale * noise
342            return vel, {
343                **state,
344                "rng_key": rng_key,
345                "istep_anneal": state["istep_anneal"] + 1,
346            }
347
348    else:
349        raise ValueError(f"Unknown thermostat {thermostat_name}")
350
351    return thermostat, postprocess, state, vel,thermostat_name
352
353
354def initialize_qtb(
355    simulation_parameters,
356    system_data,
357    fprec,
358    dt,
359    mass,
360    gamma,
361    kT,
362    species,
363    rng_key,
364    adaptive,
365    compute_thermostat_energy=False,
366):
367    state = {}
368    post_state = {}
369    qtb_parameters = simulation_parameters.get("qtb", {})
370    verbose = qtb_parameters.get("verbose", False)
371    """@keyword[fennol_md] qtb/verbose
372    Print verbose QTB thermostat information.
373    Default: False
374    """
375    if compute_thermostat_energy:
376        state["thermostat_energy"] = 0.0
377
378    mass = jnp.asarray(mass, dtype=fprec)
379
380    nat = species.shape[0]
381    # define type indices
382    species_set = set(species)
383    nspecies = len(species_set)
384    idx = {sp: i for i, sp in enumerate(species_set)}
385    type_idx = np.array([idx[sp] for sp in species], dtype=np.int32)
386    type_labels = [PERIODIC_TABLE[sp] for sp in species_set]
387
388    adapt_groups = qtb_parameters.get("adapt_groups", {})
389    """@keyword[fennol_md] qtb/adapt_groups
390    Groups of atoms with separate adQTB parameters (atoms of differents species within groups will be separated in subtypes).
391    Default: {}
392    """
393    ntypes = nspecies
394    allgroups = set()
395    for groupname, interval in adapt_groups.items():
396        indices = read_tinker_interval(interval)
397        assert allgroups.isdisjoint(indices), f"Indices in group {groupname} overlap with other groups"
398        allgroups.update(indices)
399        species_in_group = set(species[indices])
400        idx = {sp: i for i, sp in enumerate(species_in_group)}
401        for i in indices:
402            type_idx[i] = ntypes + idx[species[i]]
403        type_labels += [f"{PERIODIC_TABLE[sp]}_{groupname}" for sp in species_in_group]
404        ntypes = len(type_labels)
405
406    n_of_type = np.zeros(ntypes, dtype=np.int32)
407    for i in range(ntypes):
408        n_of_type[i] = (type_idx == i).nonzero()[0].shape[0]
409        print(f"# QTB: n({type_labels[i]}) =", n_of_type[i])
410    n_of_type = jnp.asarray(n_of_type, dtype=fprec)
411    mass_idx = jax.ops.segment_sum(mass, type_idx, ntypes) / n_of_type
412
413    niter_deconv_kin = qtb_parameters.get("niter_deconv_kin", 7)
414    """@keyword[fennol_md] qtb/niter_deconv_kin
415    Number of iterations for kinetic energy deconvolution.
416    Default: 7
417    """
418    niter_deconv_pot = qtb_parameters.get("niter_deconv_pot", 20)
419    """@keyword[fennol_md] qtb/niter_deconv_pot
420    Number of iterations for potential energy deconvolution.
421    Default: 20
422    """
423    corr_kin = qtb_parameters.get("corr_kin", -1)
424    """@keyword[fennol_md] qtb/corr_kin
425    Kinetic energy correction factor for QTB (-1 for automatic).
426    Default: -1
427    """
428    do_corr_kin = corr_kin <= 0
429    if do_corr_kin:
430        corr_kin = 1.0
431    state["corr_kin"] = corr_kin
432    post_state["corr_kin_prev"] = corr_kin
433    post_state["do_corr_kin"] = do_corr_kin
434    post_state["isame_kin"] = 0
435
436    # spectra parameters
437    omegasmear = np.pi / dt / 100.0
438    Tseg = qtb_parameters.get("tseg", 1.0 / us.PS)
439    """@keyword[fennol_md] qtb/tseg
440    Time segment length for QTB spectrum calculation.
441    Default: 1.0 ps
442    """
443    nseg = int(Tseg / dt)
444    Tseg = nseg * dt
445    dom = 2 * np.pi / (3 * Tseg)
446    omegacut = qtb_parameters.get("omegacut", 15000.0 / us.CM1)
447    """@keyword[fennol_md] qtb/omegacut
448    Cutoff frequency for QTB spectrum.
449    Default: 15000.0 cm⁻¹
450    """
451    nom = int(omegacut / dom)
452    omega = dom * np.arange((3 * nseg) // 2 + 1)
453    cutoff = jnp.asarray(
454        1.0 / (1.0 + np.exp((omega - omegacut) / omegasmear)), dtype=fprec
455    )
456    assert (
457        omegacut < omega[-1]
458    ), f"omegacut must be smaller than {omega[-1]*us.CM1} CM-1"
459
460    # initialize gammar
461    assert (
462        gamma < 0.5 * omegacut
463    ), "gamma must be much smaller than omegacut (at most 0.5*omegacut)"
464    gammar_min = qtb_parameters.get("gammar_min", 0.1)
465    """@keyword[fennol_md] qtb/gammar_min
466    Minimum value for QTB gamma ratio coefficients.
467    Default: 0.1
468    """
469    # post_state["gammar"] = jnp.asarray(np.ones((nspecies, nom)), dtype=fprec)
470    gammar = np.ones((ntypes, nom), dtype=float)
471    try:
472        for i, sp in enumerate(type_labels):
473            if not os.path.exists(f"QTB_spectra_{sp}.out"): continue
474            data = np.loadtxt(f"QTB_spectra_{sp}.out")
475            gammar[i] = data[:, 4]/(gamma*us.THZ)
476            print(f"# Restored gammar for species {sp} from QTB_spectra_{sp}.out")
477    except Exception as e:
478        print(f"# Could not restore gammar for all species with exception {e}. Starting from scratch.")
479        gammar[:,:] = 1.0
480    post_state["gammar"] = jnp.asarray(gammar, dtype=fprec)
481
482    # Ornstein-Uhlenbeck correction for colored noise
483    a1 = np.exp(-gamma * dt)
484    OUcorr = jnp.asarray(
485        (1.0 - 2.0 * a1 * np.cos(omega * dt) + a1**2) / (dt**2 * (gamma**2 + omega**2)),
486        dtype=fprec,
487    )
488
489    # hbar schedule
490    classical_kernel = qtb_parameters.get("classical_kernel", False)
491    """@keyword[fennol_md] qtb/classical_kernel
492    Use classical instead of quantum kernel for QTB.
493    Default: False
494    """
495    hbar = qtb_parameters.get("hbar", 1.0) * us.HBAR
496    """@keyword[fennol_md] qtb/hbar
497    Reduced Planck constant scaling factor for quantum effects.
498    Default: 1.0 a.u.
499    """
500    u = 0.5 * hbar * np.abs(omega) / kT
501    theta = kT * np.ones_like(omega)
502    if hbar > 0:
503        theta[1:] *= u[1:] / np.tanh(u[1:])
504    theta = jnp.asarray(theta, dtype=fprec)
505
506    noise_key, post_state["rng_key"] = jax.random.split(rng_key)
507    del rng_key
508    post_state["white_noise"] = jax.random.normal(
509        noise_key, (3 * nseg, nat, 3), dtype=jnp.float32
510    )
511
512    startsave = qtb_parameters.get("startsave", 1)
513    """@keyword[fennol_md] qtb/startsave
514    Start saving QTB statistics after this many segments.
515    Default: 1
516    """
517    counter = Counter(nseg, startsave=startsave)
518    state["istep"] = 0
519    post_state["nadapt"] = 0
520    post_state["nsample"] = 0
521
522    write_spectra = qtb_parameters.get("write_spectra", True)
523    """@keyword[fennol_md] qtb/write_spectra
524    Write QTB spectral analysis output files.
525    Default: True
526    """
527    do_compute_spectra = write_spectra or adaptive
528
529    if do_compute_spectra:
530        state["vel"] = jnp.zeros((nseg, nat, 3), dtype=fprec)
531
532        post_state["dFDT"] = jnp.zeros((ntypes, nom), dtype=fprec)
533        post_state["mCvv"] = jnp.zeros((ntypes, nom), dtype=fprec)
534        post_state["Cvf"] = jnp.zeros((ntypes, nom), dtype=fprec)
535        post_state["Cff"] = jnp.zeros((ntypes, nom), dtype=fprec)
536        post_state["dFDT_avg"] = jnp.zeros((ntypes, nom), dtype=fprec)
537        post_state["mCvv_avg"] = jnp.zeros((ntypes, nom), dtype=fprec)
538        post_state["Cvfg_avg"] = jnp.zeros((ntypes, nom), dtype=fprec)
539        post_state["Cff_avg"] = jnp.zeros((ntypes, nom), dtype=fprec)
540
541    if not adaptive:
542        update_gammar = lambda x: x
543    else:
544        # adaptation parameters
545        skipseg = qtb_parameters.get("skipseg", 1)
546        """@keyword[fennol_md] qtb/skipseg
547        Number of segments to skip before starting adaptive QTB.
548        Default: 1
549        """
550
551        adaptation_method = (
552            str(qtb_parameters.get("adaptation_method", "SIMPLE")).upper().strip()
553        )
554        """@keyword[fennol_md] qtb/adaptation_method
555        Method for adaptive QTB (SIMPLE, RATIO, ADABELIEF).
556        Default: SIMPLE
557        """
558        if adaptation_method == "SIMPLE":
559            agamma = qtb_parameters.get("agamma", 0.1)
560            """@keyword[fennol_md] qtb/agamma
561            Learning rate for adaptive QTB gamma update.
562            Default: 0.1
563            """
564            assert agamma > 0, "agamma must be positive"
565            a1_ad = agamma * Tseg  #  * gamma
566            print(f"# ADQTB SIMPLE: agamma = {agamma:.3f}")
567
568            def update_gammar(post_state):
569                g = post_state["dFDT"]
570                gammar = post_state["gammar"] - a1_ad * g
571                gammar = jnp.maximum(gammar_min, gammar)
572                return {**post_state, "gammar": gammar}
573
574        elif adaptation_method == "RATIO":
575            tau_ad = qtb_parameters.get("tau_ad", 5.0 / us.PS) 
576            """@keyword[fennol_md] qtb/tau_ad
577            Adaptation time constant for momentum averaging.
578            Default: 5.0 ps (RATIO), 1.0 ps (ADABELIEF)
579            """
580            tau_s = qtb_parameters.get("tau_s", 10 * tau_ad)
581            """@keyword[fennol_md] qtb/tau_s
582            Second moment time constant for variance averaging.
583            Default: 10*tau_ad (RATIO), 100*tau_ad (ADABELIEF)
584            """
585            assert tau_ad > 0, "tau_ad must be positive"
586            print(
587                f"# ADQTB RATIO: tau_ad = {tau_ad*us.PS:.2f} ps, tau_s = {tau_s*us.PS:.2f} ps"
588            )
589            b1 = np.exp(-Tseg / tau_ad)
590            b2 = np.exp(-Tseg / tau_s)
591            post_state["mCvv_m"] = jnp.zeros((nspecies, nom), dtype=fprec)
592            post_state["Cvf_m"] = jnp.zeros((nspecies, nom), dtype=fprec)
593            post_state["n_adabelief"] = 0
594
595            def update_gammar(post_state):
596                n_adabelief = post_state["n_adabelief"] + 1
597                mCvv_m = post_state["mCvv_m"] * b1 + post_state["mCvv"] * (1.0 - b1)
598                Cvf_m = post_state["Cvf_m"] * b2 + post_state["Cvf"] * (1.0 - b2)
599                mCvv = mCvv_m / (1.0 - b1**n_adabelief)
600                Cvf = Cvf_m / (1.0 - b2**n_adabelief)
601                # g = Cvf/(mCvv+1.e-8)-post_state["gammar"]
602                gammar = Cvf / (mCvv + 1.0e-8)
603                gammar = jnp.maximum(gammar_min, gammar)
604                return {
605                    **post_state,
606                    "gammar": gammar,
607                    "mCvv_m": mCvv_m,
608                    "Cvf_m": Cvf_m,
609                    "n_adabelief": n_adabelief,
610                }
611
612        elif adaptation_method == "ADABELIEF":
613            agamma = qtb_parameters.get("agamma", 0.1)
614            tau_ad = qtb_parameters.get("tau_ad", 1.0 / us.PS)
615            tau_s = qtb_parameters.get("tau_s", 100 * tau_ad)
616            assert tau_ad > 0, "tau_ad must be positive"
617            assert tau_s > 0, "tau_s must be positive"
618            assert agamma > 0, "agamma must be positive"
619            print(
620                f"# ADQTB ADABELIEF: agamma = {agamma:.3f}, tau_ad = {tau_ad*us.PS:.2f} ps, tau_s = {tau_s*us.PS:.2f} ps"
621            )
622
623            a1_ad = agamma * gamma  # * Tseg #* gamma
624            b1 = np.exp(-Tseg / tau_ad)
625            b2 = np.exp(-Tseg / tau_s)
626            post_state["dFDT_m"] = jnp.zeros((ntypes, nom), dtype=fprec)
627            post_state["dFDT_s"] = jnp.zeros((ntypes, nom), dtype=fprec)
628            post_state["n_adabelief"] = 0
629
630            def update_gammar(post_state):
631                n_adabelief = post_state["n_adabelief"] + 1
632                dFDT = post_state["dFDT"]
633                dFDT_m = post_state["dFDT_m"] * b1 + dFDT * (1.0 - b1)
634                dFDT_s = (
635                    post_state["dFDT_s"] * b2
636                    + (dFDT - dFDT_m) ** 2 * (1.0 - b2)
637                    + 1.0e-8
638                )
639                # bias correction
640                mt = dFDT_m / (1.0 - b1**n_adabelief)
641                st = dFDT_s / (1.0 - b2**n_adabelief)
642                gammar = post_state["gammar"] - a1_ad * mt / (st**0.5 + 1.0e-8)
643                gammar = jnp.maximum(gammar_min, gammar)
644                return {
645                    **post_state,
646                    "gammar": gammar,
647                    "dFDT_m": dFDT_m,
648                    "n_adabelief": n_adabelief,
649                    "dFDT_s": dFDT_s,
650                }
651        else: 
652            raise ValueError(
653                f"Unknown adaptation method {adaptation_method}."
654            )
655    
656    #####################
657    # RESTART
658    restart_file = system_data["name"]+".qtb.restart"
659    if os.path.exists(restart_file):
660        with open(restart_file, "rb") as f:
661            data = pickle.load(f)
662            state["corr_kin"] = data["corr_kin"]
663            post_state["corr_kin_prev"] = data["corr_kin"]
664            post_state["isame_kin"] = data["isame_kin"]
665            post_state["do_corr_kin"] = data["do_corr_kin"]
666            print(f"# Restored QTB state from {restart_file}")
667
668    def write_qtb_restart(state, post_state):
669        with open(restart_file, "wb") as f:
670            pickle.dump(
671                {
672                    "corr_kin": state["corr_kin"],
673                    "corr_kin_prev": post_state["corr_kin_prev"],
674                    "isame_kin": post_state["isame_kin"],
675                    "do_corr_kin": post_state["do_corr_kin"],
676                },
677                f,
678            )
679    ######################
680
681    def compute_corr_pot(niter=20, verbose=False):
682        if classical_kernel or hbar == 0:
683            return np.ones(nom)
684
685        s_0 = np.array((theta / kT * cutoff)[:nom])
686        s_out, s_rec, _ = deconvolute_spectrum(
687            s_0,
688            omega[:nom],
689            gamma,
690            niter,
691            kernel=kernel_lorentz_pot,
692            trans=True,
693            symmetrize=True,
694            verbose=verbose,
695        )
696        corr_pot = 1.0 + (s_out - s_0) / s_0
697        columns = np.column_stack(
698            (omega[:nom] * us.CM1, corr_pot - 1.0, s_0, s_out, s_rec)
699        )
700        np.savetxt(
701            "corr_pot.dat", columns, header="omega(cm-1) corr_pot s_0 s_out s_rec"
702        )
703        return corr_pot
704
705    def compute_corr_kin(post_state, niter=7, verbose=False):
706        if not post_state["do_corr_kin"]:
707            return post_state["corr_kin_prev"], post_state
708        if classical_kernel or hbar == 0:
709            return 1.0, post_state
710
711        K_D = post_state.get("K_D", None)
712        mCvv = (post_state["mCvv_avg"][:, :nom] * n_of_type[:, None]).sum(axis=0) / nat
713        s_0 = np.array(mCvv * kT / theta[:nom] / post_state["corr_pot"])
714        s_out, s_rec, K_D = deconvolute_spectrum(
715            s_0,
716            omega[:nom],
717            gamma,
718            niter,
719            kernel=kernel_lorentz,
720            trans=False,
721            symmetrize=True,
722            verbose=verbose,
723            K_D=K_D,
724        )
725        s_out = s_out * theta[:nom] / kT
726        s_rec = s_rec * theta[:nom] / kT * post_state["corr_pot"]
727        mCvvsum = mCvv.sum()
728        rec_ratio = mCvvsum / s_rec.sum()
729        if rec_ratio < 0.95 or rec_ratio > 1.05:
730            print(
731                f"# WARNING: reconvolution error {rec_ratio} is too high, corr_kin was not updated"
732            )
733            return post_state["corr_kin_prev"], post_state
734
735        corr_kin = mCvvsum / s_out.sum()
736        if np.abs(corr_kin - post_state["corr_kin_prev"]) < 1.0e-4:
737            isame_kin = post_state["isame_kin"] + 1
738        else:
739            isame_kin = 0
740
741        # print("# corr_kin: ", corr_kin)
742        do_corr_kin = post_state["do_corr_kin"]
743        if isame_kin > 10:
744            print(
745                "# INFO: corr_kin is converged (it did not change for 10 consecutive segments)"
746            )
747            do_corr_kin = False
748
749        return corr_kin, {
750            **post_state,
751            "corr_kin_prev": corr_kin,
752            "isame_kin": isame_kin,
753            "do_corr_kin": do_corr_kin,
754            "K_D": K_D,
755        }
756
757    @jax.jit
758    def ff_kernel(post_state):
759        if classical_kernel:
760            kernel = cutoff * (2 * gamma * kT / dt)
761        else:
762            kernel = theta * cutoff * OUcorr * (2 * gamma / dt)
763        gamma_ratio = jnp.concatenate(
764            (
765                post_state["gammar"].T * post_state["corr_pot"][:, None],
766                jnp.ones(
767                    (kernel.shape[0] - nom, ntypes), dtype=post_state["gammar"].dtype
768                ),
769            ),
770            axis=0,
771        )
772        return kernel[:, None] * gamma_ratio * mass_idx[None, :]
773
774    @jax.jit
775    def refresh_force(post_state):
776        rng_key, noise_key = jax.random.split(post_state["rng_key"])
777        white_noise = jnp.concatenate(
778            (
779                post_state["white_noise"][nseg:],
780                jax.random.normal(
781                    noise_key, (nseg, nat, 3), dtype=post_state["white_noise"].dtype
782                ),
783            ),
784            axis=0,
785        )
786        amplitude = ff_kernel(post_state) ** 0.5
787        s = jnp.fft.rfft(white_noise, 3 * nseg, axis=0) * amplitude[:, type_idx, None]
788        force = jnp.fft.irfft(s, 3 * nseg, axis=0)[nseg : 2 * nseg]
789        return force, {**post_state, "rng_key": rng_key, "white_noise": white_noise}
790
791    @jax.jit
792    def compute_spectra(force, vel, post_state):
793        sf = jnp.fft.rfft(force / gamma, 3 * nseg, axis=0, norm="ortho")
794        sv = jnp.fft.rfft(vel, 3 * nseg, axis=0, norm="ortho")
795        Cvv = jnp.sum(jnp.abs(sv[:nom]) ** 2, axis=-1).T
796        Cff = jnp.sum(jnp.abs(sf[:nom]) ** 2, axis=-1).T
797        Cvf = jnp.sum(jnp.real(sv[:nom] * jnp.conj(sf[:nom])), axis=-1).T
798
799        mCvv = (
800            (dt / 3.0)
801            * jnp.zeros_like(post_state["mCvv"]).at[type_idx].add(Cvv)
802            * mass_idx[:, None]
803            / n_of_type[:, None]
804        )
805        Cvf = (
806            (dt / 3.0)
807            * jnp.zeros_like(post_state["Cvf"]).at[type_idx].add(Cvf)
808            / n_of_type[:, None]
809        )
810        Cff = (
811            (dt / 3.0)
812            * jnp.zeros_like(post_state["Cff"]).at[type_idx].add(Cff)
813            / n_of_type[:, None]
814        )
815        dFDT = mCvv * post_state["gammar"] - Cvf
816
817        nsinv = 1.0 / post_state["nsample"]
818        b1 = 1.0 - nsinv
819        dFDT_avg = post_state["dFDT_avg"] * b1 + dFDT * nsinv
820        mCvv_avg = post_state["mCvv_avg"] * b1 + mCvv * nsinv
821        Cvfg_avg = post_state["Cvfg_avg"] * b1 + Cvf / post_state["gammar"] * nsinv
822        Cff_avg = post_state["Cff_avg"] * b1 + Cff * nsinv
823
824        return {
825            **post_state,
826            "mCvv": mCvv,
827            "Cvf": Cvf,
828            "Cff": Cff,
829            "dFDT": dFDT,
830            "dFDT_avg": dFDT_avg,
831            "mCvv_avg": mCvv_avg,
832            "Cvfg_avg": Cvfg_avg,
833            "Cff_avg": Cff_avg,
834        }
835
836    def write_spectra_to_file(post_state):
837        mCvv_avg = np.array(post_state["mCvv_avg"])
838        Cvfg_avg = np.array(post_state["Cvfg_avg"])
839        Cff_avg = np.array(post_state["Cff_avg"]) * 3.0 / dt / (gamma**2)
840        dFDT_avg = np.array(post_state["dFDT_avg"])
841        gammar = np.array(post_state["gammar"])
842        Cff_theo = np.array(ff_kernel(post_state))[:nom].T
843        for i, sp in enumerate(type_labels):
844            ff_scale = us.KELVIN / ((2 * gamma / dt) * mass_idx[i])
845            columns = np.column_stack(
846                (
847                    omega[:nom] * us.CM1,
848                    mCvv_avg[i],
849                    Cvfg_avg[i],
850                    dFDT_avg[i],
851                    gammar[i] * gamma * us.THZ,
852                    Cff_avg[i] * ff_scale,
853                    Cff_theo[i] * ff_scale,
854                )
855            )
856            np.savetxt(
857                f"QTB_spectra_{sp}.out",
858                columns,
859                fmt="%12.6f",
860                header="#omega mCvv Cvf dFDT gammar Cff",
861            )
862        if verbose:
863            print("# QTB spectra written.")
864
865    if compute_thermostat_energy:
866        state["qtb_energy_flux"] = 0.0
867
868    @jax.jit
869    def thermostat(vel, state):
870        istep = state["istep"]
871        dvel = dt * state["force"][istep] / mass[:, None]
872        new_vel = vel * a1 + dvel
873        new_state = {**state, "istep": istep + 1}
874        if do_compute_spectra:
875            vel2 = state["vel"].at[istep].set(vel * a1**0.5 + 0.5 * dvel)
876            new_state["vel"] = vel2
877        if compute_thermostat_energy:
878            dek = 0.5 * (mass[:, None] * (vel**2 - new_vel**2)).sum()
879            ekcorr = (
880                0.5
881                * (mass[:, None] * new_vel**2).sum()
882                * (1.0 - 1.0 / state.get("corr_kin", 1.0))
883            )
884            new_state["qtb_energy_flux"] = state["qtb_energy_flux"] + dek
885            new_state["thermostat_energy"] = new_state["qtb_energy_flux"] + ekcorr
886        return new_vel, new_state
887
888    @jax.jit
889    def postprocess_work(state, post_state):
890        if do_compute_spectra:
891            post_state = compute_spectra(state["force"], state["vel"], post_state)
892        if adaptive:
893            post_state = jax.lax.cond(
894                post_state["nadapt"] > skipseg, update_gammar, lambda x: x, post_state
895            )
896        new_force, post_state = refresh_force(post_state)
897        return {**state, "force": new_force}, post_state
898
899    def postprocess(state, post_state):
900        counter.increment()
901        if not counter.is_reset_step:
902            return state, post_state
903        post_state["nadapt"] += 1
904        post_state["nsample"] = max(post_state["nadapt"] - startsave + 1, 1)
905        if verbose:
906            print("# Refreshing QTB forces.")
907        state, post_state = postprocess_work(state, post_state)
908        state["corr_kin"], post_state = compute_corr_kin(post_state, niter=niter_deconv_kin)
909        state["istep"] = 0
910        if write_spectra:
911            write_spectra_to_file(post_state)
912        write_qtb_restart(state, post_state)
913        return state, post_state
914
915    post_state["corr_pot"] = jnp.asarray(compute_corr_pot(niter=niter_deconv_pot), dtype=fprec)
916
917    state["force"], post_state = refresh_force(post_state)
918    return thermostat, (postprocess, post_state), state