diff --git a/tests/test_md.py b/tests/test_md.py index dbcd5d96..9de5d539 100644 --- a/tests/test_md.py +++ b/tests/test_md.py @@ -91,6 +91,29 @@ def test_NVT_Langevin_const_T(cu_slab): assert all([at.info['MD_temperature_K'] == 500.0 for at in atoms_traj]) +def test_NVT_Langevin_const_T_per_config(cu_slab): + + calc = EMT() + + inputs = ConfigSet([cu_slab.copy(), cu_slab.copy()]) + outputs = OutputSpec() + + for at_i, at in enumerate(inputs): + at.info["WFL_MD_TEMPERATURE"] = 500 + at_i * 100 + + n_steps = 30 + + atoms_traj = md.md(inputs, outputs, calculator=calc, integrator="Langevin", steps=n_steps, dt=1.0, + temperature=200.0, temperature_tau=100/fs, rng=np.random.default_rng(1)) + + atoms_traj = list(atoms_traj) + atoms_final = atoms_traj[-1] + + assert len(atoms_traj) == (n_steps + 1) * 2 + assert all([at.info['MD_temperature_K'] == 500.0 for at in list(atoms_traj)[:n_steps + 1]]) + assert all([at.info['MD_temperature_K'] == 600.0 for at in list(atoms_traj)[n_steps + 1:]]) + + def test_NVT_const_T_mult_configs_distinct_seeds(cu_slab): calc = EMT() diff --git a/wfl/generate/md/__init__.py b/wfl/generate/md/__init__.py index 6f0399eb..fafedb90 100644 --- a/wfl/generate/md/__init__.py +++ b/wfl/generate/md/__init__.py @@ -45,12 +45,14 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBere - float: constant T - tuple/list of float, float, [int=10]: T_init, T_final, and optional number of stages for ramp - [ {'T_i': float, 'T_f' : float, 'traj_frac' : flot, 'n_stages': int=10}, ... ] list of stages, each one a ramp, with - duration defined as fraction of total number of steps + duration defined as fraction of total number of steps + Overridden by `atoms.info["WFL_MD_TEMPERATURE"]` temperature_tau: float, default None Time scale for thermostat (fs). Directly used for Berendsen integrator, or as 1/friction for Langevin integrator. pressure: None / float / tuple - applied pressure distribution (GPa) as parsed by wfl.utils.pressure.sample_pressure() - enabled Berendsen constant P volume rescaling + Applied pressure distribution (GPa) as parsed by wfl.utils.pressure.sample_pressure(). + Enables Berendsen constant P volume rescaling. + Overridden by `atoms.info["WFL_MD_PRESSURE"]` pressure_tau: float, default None time scale for Berendsen constant P volume rescaling (fs) ignored if pressure is None, defaults to 3*temperature_tau @@ -113,29 +115,42 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBere logger_constructor = logger_kwargs.pop("logger", MDLogger) logger_logfile = logger_kwargs.get("logfile", "-") - if temperature_tau is None and (temperature is not None and not isinstance(temperature, (float, int))): - raise RuntimeError(f'NVE (temperature_tau is None) can only accept temperature=float for initial T, got {type(temperature)}') - - if temperature is not None: - if isinstance(temperature, (float, int)): - # float into a list - temperature = [temperature] - if not isinstance(temperature[0], dict): - # create a stage dict from a constant or ramp - t_stage_data = temperature - # start with constant - t_stage = {'T_i': t_stage_data[0], 'T_f': t_stage_data[0], 'traj_frac': 1.0, 'n_stages': 10, 'steps': steps} - if len(t_stage_data) >= 2: - # set different final T for ramp - t_stage['T_f'] = t_stage_data[1] - if len(t_stage_data) >= 3: - # set number of stages - t_stage['n_stages'] = t_stage_data[2] - temperature = [t_stage] - else: - for t_stage in temperature: - if 'n_stages' not in t_stage: - t_stage['n_stages'] = 10 + def _get_pressure(atoms): + return atoms.info.get("WFL_MD_PRESSURE", pressure) + + def _get_temperature(atoms): + temperature_use = atoms.info.get("WFL_MD_TEMPERATURE", temperature) + + if temperature_tau is None and (temperature_use is not None and not isinstance(temperature_use, (float, int))): + raise RuntimeError(f'NVE (temperature_tau is None) can only accept temperature=float for initial T, got {type(temperature_use)}') + + if temperature_use is not None: + # assume that dicts are already in temperature profile format + if not isinstance(temperature_use, dict): + try: + # check if it's a list, tuple, etc + len(temperature_use) + except: + # number into a list + temperature_use = [temperature_use] + if not isinstance(temperature_use[0], dict): + # create a stage dict from a constant or ramp + t_stage_data = temperature_use + # start with constant + t_stage = {'T_i': t_stage_data[0], 'T_f': t_stage_data[0], 'traj_frac': 1.0, 'n_stages': 10, 'steps': steps} + if len(t_stage_data) >= 2: + # set different final T for ramp + t_stage['T_f'] = t_stage_data[1] + if len(t_stage_data) >= 3: + # set number of stages + t_stage['n_stages'] = t_stage_data[2] + temperature_use = [t_stage] + else: + for t_stage in temperature_use: + if 'n_stages' not in t_stage: + t_stage['n_stages'] = 10 + + return temperature_use for at_i, at in enumerate(atoms_to_list(atoms)): # get rng from autopara_per_item info if available ("rng" arg that was passed in was @@ -143,12 +158,15 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBere rng = _autopara_per_item_info[at_i].get("rng") item_i = _autopara_per_item_info[at_i].get("item_i") + temperature_use = _get_temperature(at) + pressure_use = _get_pressure(at) + at.calc = calculator - if pressure is not None and compressibility_au is None: - pressure = sample_pressure(pressure, at, rng=rng) + if pressure_use is not None and compressibility_au is None: + pressure_use = sample_pressure(pressure_use, at, rng=rng) at.info['MD_pressure_GPa'] = pressure # convert to ASE internal units - pressure *= GPa + pressure_use *= GPa E0 = at.get_potential_energy() c0 = at.get_cell() @@ -160,17 +178,17 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBere d2E_dF2 = (Ep + Em - 2.0 * E0) / (compressibility_fd_displ ** 2) compressibility_au = at.get_volume() / d2E_dF2 - if temperature is not None: + if temperature_use is not None: # set initial temperature assert rng is not None - MaxwellBoltzmannDistribution(at, temperature_K=temperature[0]['T_i'], force_temp=True, communicator=None, rng=rng) + MaxwellBoltzmannDistribution(at, temperature_K=temperature_use[0]['T_i'], force_temp=True, communicator=None, rng=rng) Stationary(at, preserve_temperature=True) stage_kwargs = {'timestep': dt * fs, 'logfile': logfile} if temperature_tau is None: # NVE - if pressure is not None: + if pressure_use is not None: raise RuntimeError('Cannot do NPH dynamics') md_constructor = VelocityVerlet # one stage, simple @@ -182,7 +200,7 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBere all_stage_kwargs = [] all_run_kwargs = [] - if pressure is not None: + if pressure_use is not None: md_constructor = NPTBerendsen stage_kwargs['pressure_au'] = pressure stage_kwargs['compressibility_au'] = compressibility_au @@ -199,7 +217,7 @@ def _sample_autopara_wrappable(atoms, calculator, steps, dt, integrator="NVTBere assert rng is not None stage_kwargs["rng"] = rng - for t_stage_i, t_stage in enumerate(temperature): + for t_stage_i, t_stage in enumerate(temperature_use): stage_steps = t_stage['traj_frac'] * steps if t_stage['T_f'] == t_stage['T_i']: