ai2_kit.domain.anyware module

class ai2_kit.domain.anyware.AnywareConfig(*, system_files: List[str], template_files: Mapping[str, str], product_vars: Mapping[str, List[str]] = {}, broadcast_vars: Mapping[str, List[str]] = {}, system_file_name: str, system_file_format: str, submit_script: str, post_process_fn: Optional[str] = None, delimiter: str = '$$', shuffle: bool = False)

Bases: BaseModel

broadcast_vars: Mapping[str, List[str]]

Define template variables by broadcast (broadcast as in numpy). It’s the same as product_vars, except that it will broadcast the variable to all other variables.

delimiter: str

delimiter for template

model_computed_fields: ClassVar[dict[str, ComputedFieldInfo]] = {}

A dictionary of computed field names and their corresponding ComputedFieldInfo objects.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_fields: ClassVar[dict[str, FieldInfo]] = {'broadcast_vars': FieldInfo(annotation=Mapping[str, List[str]], required=False, default={}), 'delimiter': FieldInfo(annotation=str, required=False, default='$$'), 'post_process_fn': FieldInfo(annotation=Union[str, NoneType], required=False, default=None), 'product_vars': FieldInfo(annotation=Mapping[str, List[str]], required=False, default={}), 'shuffle': FieldInfo(annotation=bool, required=False, default=False), 'submit_script': FieldInfo(annotation=str, required=True), 'system_file_format': FieldInfo(annotation=str, required=True), 'system_file_name': FieldInfo(annotation=str, required=True), 'system_files': FieldInfo(annotation=List[str], required=True), 'template_files': FieldInfo(annotation=Mapping[str, str], required=True)}

Metadata about the fields defined on the model, mapping of field names to [FieldInfo][pydantic.fields.FieldInfo].

This replaces Model.__fields__ from Pydantic V1.

post_process_fn: Optional[str]

A python function that will be executed after the task is finished. You may use this function to post-process the results.

The function must named as post_process_fn and accept a list of task directories as input. The below is an example of merging multiple file into one by keeping only the last line of each file.

post_process_fn: |

def post_process_fn(task_dirs):

import glob for task_dir in task_dirs:

files = glob.glob(os.path.join(task_dir, ‘*.out’)) # file to merge with open(os.path.join(task_dir, ‘merged.out’), ‘w’) as fp:

for file in files:

with open(file, ‘r’) as f:

lines = f.readlines() if len(lines) > 0:

fp.write(lines[-1])

product_vars: Mapping[str, List[str]]

Define template variables by Cartesian product The variable can be referenced in the template file with the following format: If there are too many variables, it will generate a large number of tasks, in this case, you can use broadcast_vars to reduce the number of tasks.

$$VAR_NAME

shuffle: bool

shuffle the combination of system_files, product_vars and broadcast_vars

submit_script: str

A bash script that will be executed in each task directory to submit the task. For example,

mpirun cp2k.popt -i cp2k.inp &> cp2k.out

system_file_format: str

The format of the system file you want to generate, for example, lammps-data, cp2k-inc, etc

For all supported data, you can refer to ase.io https://wiki.fysik.dtu.dk/ase/ase/io/io.html

Custom formats: - cp2k-inc: coord & cell in the format of CP2K include file, can be used in CP2K input file via @include coord_n_cell.inc

system_file_name: str

The name of the system file you want to generate, for example, ‘system.xyz’, ‘coord_n_cell.inc’, etc

system_files: List[str]

Artifact keys to the system

template_files: Mapping[str, str]

Templates files that will generate for each explore tasks, You can use $$VAR_NAME to reference the variables defined in product_vars and broadcast_vars.

Besides, the following build-in variables are also available: - SYSTEM_FILE: the path of the system file - DP_MODELS: the path of the deep potential models, in the format of ‘1.pb 2.pb 3.pb 4.pb’

You can use literal string to define the template file, or use !load_text to load the content from a file.

For example, if you define a template file named ‘cp2k.inp’ with the following content:

cp2k-warmup.inp: |

&GLOBAL … &END GLOBAL

cp2k.inp: !load_text cp2k.inp

class ai2_kit.domain.anyware.AnywareContext(path_prefix: str, resource_manager: ai2_kit.core.resource_manager.ResourceManager, config: ai2_kit.domain.anyware.AnywareContextConfig)

Bases: BaseCllContext

config: AnywareContextConfig

class ai2_kit.domain.anyware.AnywareContextConfig(*, script_template: BashTemplate, concurrency: int = 5)

Bases: BaseModel

concurrency: int

model_computed_fields: ClassVar[dict[str, ComputedFieldInfo]] = {}

A dictionary of computed field names and their corresponding ComputedFieldInfo objects.

model_config: ClassVar[ConfigDict] = {'extra': 'forbid'}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_fields: ClassVar[dict[str, FieldInfo]] = {'concurrency': FieldInfo(annotation=int, required=False, default=5), 'script_template': FieldInfo(annotation=BashTemplate, required=True)}

Metadata about the fields defined on the model, mapping of field names to [FieldInfo][pydantic.fields.FieldInfo].

This replaces Model.__fields__ from Pydantic V1.

script_template: BashTemplate

class ai2_kit.domain.anyware.AnywareInput(config: ai2_kit.domain.anyware.AnywareConfig, new_system_files: List[ai2_kit.core.artifact.Artifact], dp_models: Mapping[str, List[ai2_kit.core.artifact.Artifact]], type_map: List[str], mass_map: List[float])

Bases: object

config: AnywareConfig

dp_models: Mapping[str, List[Artifact]]

mass_map: List[float]

new_system_files: List[Artifact]

type_map: List[str]

class ai2_kit.domain.anyware.AnywareOutput(output_dirs: List[ai2_kit.core.artifact.Artifact])

Bases: ICllExploreOutput

get_model_devi_dataset() → List[Artifact]

output_dirs: List[Artifact]

async ai2_kit.domain.anyware.anyware(input: AnywareInput, ctx: AnywareContext) → AnywareOutput

ai2_kit.domain.anyware.make_anyware_task_dirs(work_dir: str, data_files: List[Artifact], dp_models: Mapping[str, List[str]], type_map: List[str], mass_map: List[float], product_vars: Mapping[str, List[str]], broadcast_vars: Mapping[str, List[str]], template_files: Mapping[str, str], template_delimiter: str, system_file_name: str, system_file_format: str, shuffle: bool)

ai2_kit.domain.anyware.run_post_process_fn(post_process_fn: str, task_dirs: List[str])