-
Notifications
You must be signed in to change notification settings - Fork 7
/
cll_mlp.py
343 lines (281 loc) · 12.7 KB
/
cll_mlp.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
from ai2_kit.core.executor import BaseExecutorConfig
from ai2_kit.core.artifact import ArtifactMap
from ai2_kit.core.log import get_logger
from ai2_kit.core.util import load_yaml_files, merge_dict
from ai2_kit.core.resource_manager import ResourceManager
from ai2_kit.core.checkpoint import set_checkpoint_file, apply_checkpoint
from ai2_kit.core.pydantic import BaseModel
from ai2_kit.domain import (
deepmd,
iface,
lammps,
lasp,
selector,
cp2k,
vasp,
constant as const,
updater,
)
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
from fire import Fire
import asyncio
import itertools
import copy
import os
logger = get_logger(__name__)
class CllWorkflowExecutorConfig(BaseExecutorConfig):
class Context(BaseModel):
class Train(BaseModel):
deepmd: deepmd.CllDeepmdContextConfig
class Explore(BaseModel):
lammps: Optional[lammps.CllLammpsContextConfig]
lasp: Optional[lasp.CllLaspContextConfig]
class Label(BaseModel):
cp2k: Optional[cp2k.CllCp2kContextConfig]
vasp: Optional[vasp.CllVaspContextConfig]
train: Train
explore: Explore
label: Label
context: Context
class WorkflowConfig(BaseModel):
class General(BaseModel):
type_map: List[str]
mass_map: List[float]
sel_type: Optional[List[str]]
max_iters: int = 1
mode: iface.TRAINING_MODE = 'default'
update_explore_systems: bool = False
class Label(BaseModel):
cp2k: Optional[cp2k.CllCp2kInputConfig]
vasp: Optional[vasp.CllVaspInputConfig]
class Train(BaseModel):
deepmd: deepmd.CllDeepmdInputConfig
class Explore(BaseModel):
lammps: Optional[lammps.CllLammpsInputConfig]
lasp: Optional[lasp.CllLaspInputConfig]
class Select(BaseModel):
model_devi: selector.CllModelDeviSelectorInputConfig
class Update(BaseModel):
walkthrough: updater.CllWalkthroughUpdaterInputConfig
general: General
train: Train
explore: Explore
select: Select
label: Label
update: Update
class CllWorkflowConfig(BaseModel):
executors: Dict[str, CllWorkflowExecutorConfig]
artifacts: ArtifactMap
workflow: Any # Keep it raw here, it should be parsed later in iteration
def run_workflow(*config_files,
executor: Optional[str] = None,
path_prefix: Optional[str] = None,
checkpoint: Optional[str] = None):
"""
Run Closed-Loop Learning (CLL) workflow to train Machine Learning Potential (MLP) models.
Args:
config_files: path of config files, should be yaml files, can be multiple, support glob pattern
executor: name of executor, should be defined in config `executors` section
path_prefix: path prefix for output
checkpoint: checkpoint file
"""
if checkpoint is not None:
set_checkpoint_file(checkpoint)
config_data = load_yaml_files(*config_files)
config = CllWorkflowConfig.parse_obj(config_data)
if executor not in config.executors:
raise ValueError(f'executor {executor} is not found')
if path_prefix is None:
raise ValueError('path_prefix should not be empty')
iface.init_artifacts(config.artifacts)
resource_manager = ResourceManager(
executor_configs=config.executors,
artifacts=config.artifacts,
default_executor=executor,
)
return asyncio.run(cll_mlp_training_workflow(config, resource_manager, executor, path_prefix))
async def cll_mlp_training_workflow(config: CllWorkflowConfig,
resource_manager: ResourceManager,
executor: str,
path_prefix: str):
context_config = config.executors[executor].context
raw_workflow_config = copy.deepcopy(config.workflow)
# output of each step
label_output: Optional[iface.ICllLabelOutput] = None
selector_output: Optional[iface.ICllSelectorOutput] = None
train_output: Optional[iface.ICllTrainOutput] = None
explore_output: Optional[iface.ICllExploreOutput] = None
# cursor of update table
update_cursor = 0
# Start iteration
for i in itertools.count(0):
# parse workflow config
workflow_config = WorkflowConfig.parse_obj(raw_workflow_config)
shared_vars = precondition(workflow_config)
if i >= workflow_config.general.max_iters:
logger.info(f'Iteration {i} exceeds max_iters, stop iteration.')
break
# shortcut for type_map and mass_map
type_map = workflow_config.general.type_map
mass_map = workflow_config.general.mass_map
# decide path prefix for each iteration
iter_path_prefix = os.path.join(path_prefix, f'iters-{i:03d}')
# prefix of checkpoint
cp_prefix = f'iters-{i:03d}'
# label
if workflow_config.label.cp2k and context_config.label.cp2k:
cp2k_input = cp2k.CllCp2kInput(
config=workflow_config.label.cp2k,
mode=workflow_config.general.mode,
type_map=type_map,
system_files=[] if selector_output is None else selector_output.get_model_devi_dataset(),
initiated=i > 0,
)
cp2k_context = cp2k.CllCp2kContext(
config=context_config.label.cp2k,
path_prefix=os.path.join(iter_path_prefix, 'label-cp2k'),
resource_manager=resource_manager,
)
label_output = await apply_checkpoint(f'{cp_prefix}/label-cp2k')(cp2k.cll_cp2k)(cp2k_input, cp2k_context)
elif workflow_config.label.vasp and context_config.label.vasp:
vasp_input = vasp.CllVaspInput(
config=workflow_config.label.vasp,
type_map=type_map,
system_files=[] if selector_output is None else selector_output.get_model_devi_dataset(),
initiated=i > 0,
)
vasp_context = vasp.CllVaspContext(
config=context_config.label.vasp,
path_prefix=os.path.join(iter_path_prefix, 'label-vasp'),
resource_manager=resource_manager,
)
label_output = await apply_checkpoint(f'{cp_prefix}/label-vasp')(vasp.cll_vasp)(vasp_input, vasp_context)
else:
raise ValueError('No label method is specified')
# return if no new data is generated
if i > 0 and len(label_output.get_labeled_system_dataset()) == 0:
logger.info("No new data is generated, stop iteration.")
break
# train
if workflow_config.train.deepmd:
deepmd_input = deepmd.CllDeepmdInput(
config=workflow_config.train.deepmd,
mode=workflow_config.general.mode,
type_map=type_map,
old_dataset=[] if train_output is None else train_output.get_training_dataset(),
new_dataset=label_output.get_labeled_system_dataset(),
sel_type=shared_vars.dp_sel_type,
)
deepmd_context = deepmd.CllDeepmdContext(
path_prefix=os.path.join(iter_path_prefix, 'train-deepmd'),
config=context_config.train.deepmd,
resource_manager=resource_manager,
)
train_output = await apply_checkpoint(f'{cp_prefix}/train-deepmd')(deepmd.cll_deepmd)(deepmd_input, deepmd_context)
else:
raise ValueError('No train method is specified')
# explore
new_explore_system_files = []
if workflow_config.general.update_explore_systems and selector_output is not None:
new_explore_system_files = selector_output.get_new_explore_systems()
if workflow_config.explore.lammps and context_config.explore.lammps:
lammps_input = lammps.CllLammpsInput(
config=workflow_config.explore.lammps,
mode=workflow_config.general.mode,
type_map=type_map,
mass_map=mass_map,
dp_models={'': train_output.get_mlp_models()},
preset_template='default',
new_system_files=new_explore_system_files,
dp_modifier=shared_vars.dp_modifier,
dp_sel_type=shared_vars.dp_sel_type,
)
lammps_context = lammps.CllLammpsContext(
path_prefix=os.path.join(iter_path_prefix, 'explore-lammps'),
config=context_config.explore.lammps,
resource_manager=resource_manager,
)
explore_output = await apply_checkpoint(f'{cp_prefix}/explore-lammps')(lammps.cll_lammps)(lammps_input, lammps_context)
elif workflow_config.explore.lasp and context_config.explore.lasp:
lasp_input = lasp.CllLaspInput(
config=workflow_config.explore.lasp,
type_map=type_map,
mass_map=mass_map,
models=train_output.get_mlp_models(),
new_system_files=new_explore_system_files,
)
lasp_context = lasp.CllLaspContext(
config=context_config.explore.lasp,
path_prefix=os.path.join(iter_path_prefix, 'explore-lasp'),
resource_manager=resource_manager,
)
explore_output = await apply_checkpoint(f'{cp_prefix}/explore-lasp')(lasp.cll_lasp)(lasp_input, lasp_context)
else:
raise ValueError('No explore method is specified')
# select
if workflow_config.select.model_devi:
selector_input = selector.CllModelDeviSelectorInput(
config=workflow_config.select.model_devi,
model_devi_data=explore_output.get_model_devi_dataset(),
model_devi_file=const.MODEL_DEVI_OUT,
type_map=type_map,
)
selector_context = selector.CllModelDevSelectorContext(
path_prefix=os.path.join(iter_path_prefix, 'selector-model-devi'),
resource_manager=resource_manager,
)
selector_output = await apply_checkpoint(f'{cp_prefix}/selector-model-devi')(selector.cll_model_devi_selector)(selector_input, selector_context)
else:
raise ValueError('No select method is specified')
# Update
update_config = workflow_config.update.walkthrough
# nothing to update because the table is empty
if not update_config.table:
continue
# keep using the latest config when it reach the end of table
if update_cursor >= len(update_config.table):
continue
# move cursor to next row if passing rate pass threshold
if selector_output.get_passing_rate() > update_config.passing_rate_threshold:
raw_workflow_config = merge_dict(copy.deepcopy(
config.workflow), update_config.table[update_cursor])
update_cursor += 1
@dataclass
class SharedVars():
dp_modifier: Optional[dict] = None
dp_sel_type: Optional[List[int]] = None
def precondition(workflow_cfg: WorkflowConfig) -> SharedVars:
"""
precondition of workflow config to raise error early,
and extra variables that may shared by multiple steps
The known shared variables are:
dp_modifier, which include vars sys_charge_map, model_charge_map, ewald_h, ewald_beta
sel_type, which is suppose to be used in dplr/dpff mode
"""
shared_vars = SharedVars()
mode = workflow_cfg.general.mode
type_map = workflow_cfg.general.type_map
sel_type = workflow_cfg.general.sel_type
if mode == 'dpff':
assert sel_type is not None, 'sel_type should be specified in general config for dpff mode'
shared_vars.dp_sel_type = [ type_map.index(t) for t in sel_type ]
deepmd_cfg = workflow_cfg.train.deepmd
if deepmd_cfg is not None:
if mode == 'dpff':
modifier = deepmd_cfg.input_template['model'].get('modifier')
assert modifier is not None, 'modifier should be specified in deepmd input template for dpff mode'
shared_vars.dp_modifier = modifier
lammps_cfg = workflow_cfg.explore.lammps
if lammps_cfg is not None:
if mode == 'dpff':
lammps_cfg.assert_var('EFIELD')
lammps_cfg.assert_var('KMESH')
efield = lammps_cfg.explore_vars.get('EFIELD', lammps_cfg.broadcast_vars.get('EFIELD'))
assert all([isinstance(item, list) for item in efield ]), 'EFIELD should be a list of vector' # type: ignore
elif mode in ['fep-redox', 'fep-pka']:
lammps_cfg.assert_var('LAMBDA_f')
return shared_vars
if __name__ == '__main__':
# use python-fire to parse command line arguments
Fire(run_workflow)