fennol.md.dynamic

  1import sys, os, io
  2import argparse
  3import time
  4from pathlib import Path
  5
  6import numpy as np
  7from collections import defaultdict
  8import jax
  9import jax.numpy as jnp
 10import pickle
 11import yaml
 12
 13from ..utils.io import (
 14    write_arc_frame,
 15    write_xyz_frame,
 16    write_extxyz_frame,
 17    human_time_duration,
 18)
 19from .utils import wrapbox, save_dynamics_restart
 20from ..utils import minmaxone, AtomicUnits as au,read_tinker_interval
 21from ..utils.input_parser import parse_input,convert_dict_units, InputFile
 22from .integrate import initialize_dynamics
 23
 24
 25def main():
 26    # os.environ["OMP_NUM_THREADS"] = "1"
 27    sys.stdout = io.TextIOWrapper(
 28        open(sys.stdout.fileno(), "wb", 0), write_through=True
 29    )
 30    ### Read the parameter file
 31    parser = argparse.ArgumentParser(prog="fennol_md")
 32    parser.add_argument("param_file", type=Path, help="Parameter file")
 33    args = parser.parse_args()
 34
 35    param_file = args.param_file
 36    if not param_file.exists() and not param_file.is_file():
 37        raise FileNotFoundError(f"Parameter file {param_file} not found")
 38    
 39    if param_file.suffix in [".yaml", ".yml"]:
 40        with open(param_file, "r") as f:
 41            simulation_parameters = convert_dict_units(yaml.safe_load(f))
 42            simulation_parameters = InputFile(**simulation_parameters)
 43    elif param_file.suffix == ".fnl":
 44        simulation_parameters = parse_input(args.param_file)
 45    else:
 46        raise ValueError(
 47            f"Unknown parameter file format '{param_file.suffix}'. Supported formats are '.yaml', '.yml' and '.fnl'"
 48        )
 49
 50    ### Set the device
 51    if "FENNOL_DEVICE" in os.environ:
 52        device = os.environ["FENNOL_DEVICE"].lower()
 53        print(f"# Setting device from env FENNOL_DEVICE={device}")
 54    else:
 55        device = simulation_parameters.get("device", "cpu").lower()
 56    if device == "cpu":
 57        jax.config.update('jax_platforms', 'cpu')
 58        os.environ["CUDA_VISIBLE_DEVICES"] = ""
 59    elif device.startswith("cuda") or device.startswith("gpu"):
 60        if ":" in device:
 61            num = device.split(":")[-1]
 62            os.environ["CUDA_VISIBLE_DEVICES"] = num
 63        else:
 64            os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 65        device = "gpu"
 66
 67    _device = jax.devices(device)[0]
 68    jax.config.update("jax_default_device", _device)
 69
 70    ### Set the precision
 71    enable_x64 = simulation_parameters.get("double_precision", False)
 72    jax.config.update("jax_enable_x64", enable_x64)
 73    fprec = "float64" if enable_x64 else "float32"
 74
 75    matmul_precision = simulation_parameters.get("matmul_prec", "highest").lower()
 76    assert matmul_precision in [
 77        "default",
 78        "high",
 79        "highest",
 80    ], "matmul_prec must be one of 'default','high','highest'"
 81    if matmul_precision != "highest":
 82        print(f"# Setting matmul precision to '{matmul_precision}'")
 83    if matmul_precision == "default" and fprec == "float32":
 84        print(
 85            "# Warning: default matmul precision involves float16 operations which may lead to large numerical errors on energy and pressure estimations ! It is recommended to set matmul_prec to 'high' or 'highest'."
 86        )
 87    jax.config.update("jax_default_matmul_precision", matmul_precision)
 88
 89    # with jax.default_device(_device):
 90    dynamic(simulation_parameters, device, fprec)
 91
 92
 93def dynamic(simulation_parameters, device, fprec):
 94    tstart_dyn = time.time()
 95
 96    random_seed = simulation_parameters.get(
 97        "random_seed", np.random.randint(0, 2**32 - 1)
 98    )
 99    print(f"# random_seed: {random_seed}")
100    rng_key = jax.random.PRNGKey(random_seed)
101
102    ## INITIALIZE INTEGRATOR AND SYSTEM
103    rng_key, subkey = jax.random.split(rng_key)
104    step, update_conformation, system_data, dyn_state, conformation, system = (
105        initialize_dynamics(simulation_parameters, fprec, subkey)
106    )
107
108    nat = system_data["nat"]
109    dt = dyn_state["dt"]
110    ## get number of steps
111    nsteps = int(simulation_parameters.get("nsteps"))
112    start_time_ps = dyn_state.get("start_time_ps", 0.0)
113
114    ### Set I/O parameters
115    Tdump = simulation_parameters.get("tdump", 1.0 / au.PS) * au.FS
116    ndump = int(Tdump / dt)
117    system_name = system_data["name"]
118    estimate_pressure = dyn_state["estimate_pressure"]
119
120    @jax.jit
121    def check_nan(system):
122        return jnp.any(jnp.isnan(system["vel"])) | jnp.any(
123            jnp.isnan(system["coordinates"])
124        )
125
126    if system_data["pbc"] is not None:
127        cell = system["cell"]
128        reciprocal_cell = np.linalg.inv(cell)
129        do_wrap_box = simulation_parameters.get("wrap_box", False)
130        if do_wrap_box:
131            wrap_groups_def = simulation_parameters.get("wrap_groups",None)
132            if wrap_groups_def is None:
133                wrap_groups = None
134            else:
135                wrap_groups = {}
136                assert isinstance(wrap_groups_def, dict), "wrap_groups must be a dictionary"
137                for k, v in wrap_groups_def.items():
138                    wrap_groups[k]=read_tinker_interval(v)
139                # check that pairwise intersection of wrap groups is empty
140                wrap_groups_keys = list(wrap_groups.keys())
141                for i in range(len(wrap_groups_keys)):
142                    i_key = wrap_groups_keys[i]
143                    w1 = set(wrap_groups[i_key])
144                    for j in range(i + 1, len(wrap_groups_keys)):
145                        j_key = wrap_groups_keys[j]
146                        w2 = set(wrap_groups[j_key])
147                        if  w1.intersection(w2):
148                            raise ValueError(
149                                f"Wrap groups {i_key} and {j_key} have common atoms: {w1.intersection(w2)}"
150                            )
151                group_all = np.concatenate(list(wrap_groups.values()))
152                # get all atoms that are not in any wrap group
153                group_none = np.setdiff1d(np.arange(nat), group_all)
154                print(f"# Wrap groups: {wrap_groups}")
155                wrap_groups["__other"] = group_none
156                wrap_groups = ((k, v) for k, v in wrap_groups.items())
157
158    else:
159        cell = None
160        reciprocal_cell = None
161        do_wrap_box = False
162        wrap_groups = None
163
164    ### Energy units and print initial energy
165    model_energy_unit = system_data["model_energy_unit"]
166    model_energy_unit_str = system_data["model_energy_unit_str"]
167    per_atom_energy = simulation_parameters.get("per_atom_energy", True)
168    energy_unit_str = system_data["energy_unit_str"]
169    energy_unit = system_data["energy_unit"]
170    print("# Energy unit: ", energy_unit_str)
171    atom_energy_unit = energy_unit
172    atom_energy_unit_str = energy_unit_str
173    if per_atom_energy:
174        atom_energy_unit /= nat
175        atom_energy_unit_str = f"{energy_unit_str}/atom"
176        print("# Printing Energy per atom")
177    print(
178        f"# Initial potential energy: {system['epot']*atom_energy_unit}; kinetic energy: {system['ek']*atom_energy_unit}"
179    )
180    f = system["forces"]
181    minmaxone(jnp.abs(f * energy_unit), "# forces min/max/rms:")
182
183    ## printing options
184    print_timings = simulation_parameters.get("print_timings", False)
185    nprint = int(simulation_parameters.get("nprint", 10))
186    assert nprint > 0, "nprint must be > 0"
187    nsummary = simulation_parameters.get("nsummary", 100 * nprint)
188    assert nsummary > nprint, "nsummary must be > nprint"
189
190    save_keys = simulation_parameters.get("save_keys", [])
191    if save_keys:
192        print(f"# Saving keys: {save_keys}")
193        fkeys = open(f"{system_name}.traj.pkl", "wb+")
194    else:
195        fkeys = None
196    
197    ### initialize colvars
198    use_colvars = "colvars" in dyn_state
199    if use_colvars:
200        print(f"# Colvars: {dyn_state['colvars']}")
201        colvars_names = dyn_state["colvars"]
202        # open colvars file and print header
203        fcolvars = open(f"{system_name}.colvars.traj", "a")
204        fcolvars.write("#time[ps] ")
205        fcolvars.write(" ".join(colvars_names))
206        fcolvars.write("\n")
207        fcolvars.flush()
208
209    ### Print header
210    include_thermostat_energy = "thermostat_energy" in system["thermostat"]
211    thermostat_name = dyn_state["thermostat_name"]
212    pimd = dyn_state["pimd"]
213    variable_cell = dyn_state["variable_cell"]
214    nbeads = system_data.get("nbeads", 1)
215    dyn_name = "PIMD" if pimd else "MD"
216    print("#" * 84)
217    print(
218        f"# Running {nsteps:_} steps of {thermostat_name} {dyn_name} simulation on {device}"
219    )
220    header = "#     Step   Time[ps]        Etot        Epot        Ekin    Temp[K]"
221    if pimd:
222        header += "  Temp_c[K]"
223    if include_thermostat_energy:
224        header += "      Etherm"
225    if estimate_pressure:
226        print_aniso_pressure = simulation_parameters.get("print_aniso_pressure", False)
227        pressure_unit_str = simulation_parameters.get("pressure_unit", "atm")
228        pressure_unit = au.get_multiplier(pressure_unit_str) * au.BOHR**3
229        header += f"  Press[{pressure_unit_str}]"
230    if variable_cell:
231        header += "   Density"
232    print(header)
233
234    ### Open trajectory file
235    traj_format = simulation_parameters.get("traj_format", "arc").lower()
236    if traj_format == "xyz":
237        traj_ext = ".traj.xyz"
238        write_frame = write_xyz_frame
239    elif traj_format == "extxyz":
240        traj_ext = ".traj.extxyz"
241        write_frame = write_extxyz_frame
242    elif traj_format == "arc":
243        traj_ext = ".arc"
244        write_frame = write_arc_frame
245    else:
246        raise ValueError(
247            f"Unknown trajectory format '{traj_format}'. Supported formats are 'arc' and 'xyz'"
248        )
249
250    write_all_beads = simulation_parameters.get("write_all_beads", False) and pimd
251
252    if write_all_beads:
253        fout = [
254            open(f"{system_name}_bead{i+1:03d}" + traj_ext, "a") for i in range(nbeads)
255        ]
256    else:
257        fout = open(system_name + traj_ext, "a")
258
259    ensemble_key = simulation_parameters.get("etot_ensemble_key", None)
260    if ensemble_key is not None:
261        fens = open(f"{system_name}.ensemble_weights.traj", "a")
262
263    write_centroid = simulation_parameters.get("write_centroid", False) and pimd
264    if write_centroid:
265        fcentroid = open(f"{system_name}_centroid" + traj_ext, "a")
266
267    ### initialize proprerty trajectories
268    properties_traj = defaultdict(list)
269
270    ### initialize counters and timers
271    t0 = time.time()
272    t0dump = t0
273    istep = 0
274    t0full = time.time()
275    force_preprocess = False
276
277    for istep in range(1, nsteps + 1):
278
279        ### update the system
280        dyn_state, system, conformation, model_output = step(
281            istep, dyn_state, system, conformation, force_preprocess
282        )
283
284        ### print properties
285        if istep % nprint == 0:
286            t1 = time.time()
287            tperstep = (t1 - t0) / nprint
288            t0 = t1
289            nsperday = (24 * 60 * 60 / tperstep) * dt / 1e6
290
291            ek = system["ek"]
292            epot = system["epot"]
293            etot = ek + epot
294            temper = 2 * ek / (3.0 * nat) * au.KELVIN
295
296            th_state = system["thermostat"]
297            if include_thermostat_energy:
298                etherm = th_state["thermostat_energy"]
299                etot = etot + etherm
300
301            properties_traj[f"Etot[{atom_energy_unit_str}]"].append(
302                etot * atom_energy_unit
303            )
304            properties_traj[f"Epot[{atom_energy_unit_str}]"].append(
305                epot * atom_energy_unit
306            )
307            properties_traj[f"Ekin[{atom_energy_unit_str}]"].append(
308                ek * atom_energy_unit
309            )
310            properties_traj["Temper[Kelvin]"].append(temper)
311            if pimd:
312                ek_c = system["ek_c"]
313                temper_c = 2 * ek_c / (3.0 * nat) * au.KELVIN
314                properties_traj["Temper_c[Kelvin]"].append(temper_c)
315
316            ### construct line of properties
317            simulated_time = start_time_ps + istep * dt / 1000
318            line = f"{istep:10.6g} {simulated_time: 10.3f}  {etot*atom_energy_unit: #10.4f}  {epot*atom_energy_unit: #10.4f}  {ek*atom_energy_unit: #10.4f} {temper: 10.2f}"
319            if pimd:
320                line += f" {temper_c: 10.2f}"
321            if include_thermostat_energy:
322                line += f"  {etherm*atom_energy_unit: #10.4f}"
323                properties_traj[f"Etherm[{atom_energy_unit_str}]"].append(
324                    etherm * atom_energy_unit
325                )
326            if estimate_pressure:
327                pres = system["pressure"] * pressure_unit
328                properties_traj[f"Pressure[{pressure_unit_str}]"].append(pres)
329                if print_aniso_pressure:
330                    pres_tensor = system["pressure_tensor"] * pressure_unit
331                    pres_tensor = 0.5 * (pres_tensor + pres_tensor.T)
332                    properties_traj[f"Pressure_xx[{pressure_unit_str}]"].append(
333                        pres_tensor[0, 0]
334                    )
335                    properties_traj[f"Pressure_yy[{pressure_unit_str}]"].append(
336                        pres_tensor[1, 1]
337                    )
338                    properties_traj[f"Pressure_zz[{pressure_unit_str}]"].append(
339                        pres_tensor[2, 2]
340                    )
341                    properties_traj[f"Pressure_xy[{pressure_unit_str}]"].append(
342                        pres_tensor[0, 1]
343                    )
344                    properties_traj[f"Pressure_xz[{pressure_unit_str}]"].append(
345                        pres_tensor[0, 2]
346                    )
347                    properties_traj[f"Pressure_yz[{pressure_unit_str}]"].append(
348                        pres_tensor[1, 2]
349                    )
350                line += f" {pres:10.3f}"
351            if variable_cell:
352                density = system["density"]
353                properties_traj["Density[g/cm^3]"].append(density)
354                if print_aniso_pressure:
355                    cell = system["cell"]
356                    properties_traj[f"Cell_Ax[Angstrom]"].append(cell[0, 0])
357                    properties_traj[f"Cell_Ay[Angstrom]"].append(cell[0, 1])
358                    properties_traj[f"Cell_Az[Angstrom]"].append(cell[0, 2])
359                    properties_traj[f"Cell_Bx[Angstrom]"].append(cell[1, 0])
360                    properties_traj[f"Cell_By[Angstrom]"].append(cell[1, 1])
361                    properties_traj[f"Cell_Bz[Angstrom]"].append(cell[1, 2])
362                    properties_traj[f"Cell_Cx[Angstrom]"].append(cell[2, 0])
363                    properties_traj[f"Cell_Cy[Angstrom]"].append(cell[2, 1])
364                    properties_traj[f"Cell_Cz[Angstrom]"].append(cell[2, 2])
365                line += f" {density:10.4f}"
366                if "piston_temperature" in system["barostat"]:
367                    piston_temperature = system["barostat"]["piston_temperature"]
368                    properties_traj["T_piston[Kelvin]"].append(piston_temperature)
369
370            print(line)
371
372            ### save colvars
373            if use_colvars:
374                colvars = system["colvars"]
375                fcolvars.write(f"{simulated_time:.3f} ")
376                fcolvars.write(" ".join([f"{colvars[k]:.6f}" for k in colvars_names]))
377                fcolvars.write("\n")
378                fcolvars.flush()
379
380            if save_keys:
381                data = {
382                    k: (
383                        np.asarray(model_output[k])
384                        if isinstance(model_output[k], jnp.ndarray)
385                        else model_output[k]
386                    )
387                    for k in save_keys
388                }
389                data["properties"] = {
390                    k: float(v) for k, v in zip(header.split()[1:], line.split())
391                }
392                data["properties"]["properties_energy_unit"] = (atom_energy_unit,atom_energy_unit_str)
393                data["properties"]["model_energy_unit"] = (model_energy_unit,model_energy_unit_str)
394
395                pickle.dump(data, fkeys)
396
397        ### save frame
398        if istep % ndump == 0:
399            line = "# Write XYZ frame"
400            if variable_cell:
401                cell = np.array(system["cell"])
402                reciprocal_cell = np.linalg.inv(cell)
403            if do_wrap_box:
404                if pimd:
405                    centroid = wrapbox(system["coordinates"][0], cell, reciprocal_cell,wrap_groups=wrap_groups)
406                    system["coordinates"] = system["coordinates"].at[0].set(centroid)
407                else:
408                    system["coordinates"] = wrapbox(
409                        system["coordinates"], cell, reciprocal_cell,wrap_groups=wrap_groups
410                    )
411                conformation = update_conformation(conformation, system)
412                line += " (atoms have been wrapped into the box)"
413                force_preprocess = True
414            print(line)
415
416            save_dynamics_restart(system_data, conformation, dyn_state, system)
417
418            properties = {
419                "energy": float(system["epot"]) * energy_unit,
420                "Time_ps": start_time_ps + istep * dt / 1000,
421                "energy_unit": energy_unit_str,
422            }
423
424            if write_all_beads:
425                coords = np.asarray(conformation["coordinates"].reshape(-1, nat, 3))
426                for i, fb in enumerate(fout):
427                    write_frame(
428                        fb,
429                        system_data["symbols"],
430                        coords[i],
431                        cell=cell,
432                        properties=properties,
433                        forces=None,  # np.asarray(system["forces"].reshape(nbeads, nat, 3)[0]) * energy_unit,
434                    )
435            else:
436                write_frame(
437                    fout,
438                    system_data["symbols"],
439                    np.asarray(conformation["coordinates"].reshape(-1, nat, 3)[0]),
440                    cell=cell,
441                    properties=properties,
442                    forces=None,  # np.asarray(system["forces"].reshape(nbeads, nat, 3)[0]) * energy_unit,
443                )
444            if write_centroid:
445                centroid = np.asarray(system["coordinates"][0])
446                write_frame(
447                    fcentroid,
448                    system_data["symbols"],
449                    centroid,
450                    cell=cell,
451                    properties=properties,
452                    forces=np.asarray(system["forces"].reshape(nbeads, nat, 3)[0])
453                    * energy_unit,
454                )
455            if ensemble_key is not None:
456                weights = " ".join(
457                    [f"{w:.6f}" for w in system["ensemble_weights"].tolist()]
458                )
459                fens.write(f"{weights}\n")
460                fens.flush()
461
462        ### summary over last nsummary steps
463        if istep % (nsummary) == 0:
464            if check_nan(system):
465                raise ValueError(f"dynamics crashed at step {istep}.")
466            tfull = time.time() - t0full
467            t0full = time.time()
468            tperstep = tfull / (nsummary)
469            nsperday = (24 * 60 * 60 / tperstep) * dt / 1e6
470            elapsed_time = time.time() - tstart_dyn
471            estimated_remaining_time = tperstep * (nsteps - istep)
472            estimated_total_time = elapsed_time + estimated_remaining_time
473
474            print("#" * 50)
475            print(f"# Step {istep:_} of {nsteps:_}  ({istep/nsteps*100:.5g} %)")
476            print(f"# Simulated time      : {istep * dt*1.e-3:.3f} ps")
477            print(f"# Tot. Simu. time     : {start_time_ps + istep * dt*1.e-3:.3f} ps")
478            print(f"# Tot. elapsed time   : {human_time_duration(elapsed_time)}")
479            print(
480                f"# Est. total duration   : {human_time_duration(estimated_total_time)}"
481            )
482            print(
483                f"# Est. remaining time : {human_time_duration(estimated_remaining_time)}"
484            )
485            print(f"# Time for {nsummary:_} steps : {human_time_duration(tfull)}")
486
487            corr_kin = system["thermostat"].get("corr_kin", None)
488            if corr_kin is not None:
489                print(f"# QTB kin. correction : {100*(corr_kin-1.):.2f} %")
490            print(f"# Averages over last {nsummary:_} steps :")
491            for k, v in properties_traj.items():
492                if len(properties_traj[k]) == 0:
493                    continue
494                mu = np.mean(properties_traj[k])
495                sig = np.std(properties_traj[k])
496                ksplit = k.split("[")
497                name = ksplit[0].strip()
498                unit = ksplit[1].replace("]", "").strip() if len(ksplit) > 1 else ""
499                print(f"#   {name:10} : {mu: #10.5g}   +/- {sig: #9.3g}  {unit}")
500
501            print(f"# Perf.: {nsperday:.2f} ns/day  ( {1.0 / tperstep:.2f} step/s )")
502            print("#" * 50)
503            if istep < nsteps:
504                print(header)
505            ## reset property trajectories
506            properties_traj = defaultdict(list)
507
508    print(f"# Run done in {human_time_duration(time.time()-tstart_dyn)}")
509    ### close trajectory file
510    fout.close()
511    if ensemble_key is not None:
512        fens.close()
513    if use_colvars:
514        fcolvars.close()
515    if fkeys is not None:
516        fkeys.close()
517    if write_centroid:
518        fcentroid.close()
519
520
521if __name__ == "__main__":
522    main()