ai2_kit.domain.anyware module#

pydantic model ai2_kit.domain.anyware.AnywareConfig[source]#

Bases: BaseModel

Show JSON schema

{
   "title": "AnywareConfig",
   "type": "object",
   "properties": {
      "system_files": {
         "items": {
            "type": "string"
         },
         "title": "System Files",
         "type": "array"
      },
      "template_files": {
         "additionalProperties": {
            "type": "string"
         },
         "title": "Template Files",
         "type": "object"
      },
      "product_vars": {
         "additionalProperties": {
            "items": {
               "type": "string"
            },
            "type": "array"
         },
         "default": {},
         "title": "Product Vars",
         "type": "object"
      },
      "broadcast_vars": {
         "additionalProperties": {
            "items": {
               "type": "string"
            },
            "type": "array"
         },
         "default": {},
         "title": "Broadcast Vars",
         "type": "object"
      },
      "system_file_name": {
         "title": "System File Name",
         "type": "string"
      },
      "system_file_format": {
         "title": "System File Format",
         "type": "string"
      },
      "submit_script": {
         "title": "Submit Script",
         "type": "string"
      },
      "post_process_fn": {
         "anyOf": [
            {
               "type": "string"
            },
            {
               "type": "null"
            }
         ],
         "default": null,
         "title": "Post Process Fn"
      },
      "delimiter": {
         "default": "$$",
         "title": "Delimiter",
         "type": "string"
      },
      "shuffle": {
         "default": false,
         "title": "Shuffle",
         "type": "boolean"
      }
   },
   "additionalProperties": false,
   "required": [
      "system_files",
      "template_files",
      "system_file_name",
      "system_file_format",
      "submit_script"
   ]
}

Config:

extra: str = forbid

Fields:

broadcast_vars (Mapping[str, List[str]])
delimiter (str)
post_process_fn (str | None)
product_vars (Mapping[str, List[str]])
shuffle (bool)
submit_script (str)
system_file_format (str)
system_file_name (str)
system_files (List[str])
template_files (Mapping[str, str])

field 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.

field delimiter: str = '$$'#: delimiter for template

field post_process_fn: str | None = None#

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])

field 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

field shuffle: bool = False#: shuffle the combination of system_files, product_vars and broadcast_vars

field submit_script: str [Required]#

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

field system_file_format: str [Required]#

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

field system_file_name: str [Required]#: The name of the system file you want to generate, for example, ‘system.xyz’, ‘coord_n_cell.inc’, etc

field system_files: List[str] [Required]#: Artifact keys to the system

field template_files: Mapping[str, str] [Required]#

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)[source]#

Bases: BaseCllContext

config: AnywareContextConfig#

pydantic model ai2_kit.domain.anyware.AnywareContextConfig[source]#

Bases: BaseModel

Show JSON schema

{
   "title": "AnywareContextConfig",
   "type": "object",
   "properties": {
      "script_template": {
         "$ref": "#/$defs/BashTemplate"
      },
      "concurrency": {
         "default": 5,
         "title": "Concurrency",
         "type": "integer"
      }
   },
   "$defs": {
      "BashTemplate": {
         "additionalProperties": false,
         "properties": {
            "shebang": {
               "default": "#!/bin/bash",
               "title": "Shebang",
               "type": "string"
            },
            "header": {
               "default": "",
               "title": "Header",
               "type": "string"
            },
            "setup": {
               "default": "",
               "title": "Setup",
               "type": "string"
            },
            "teardown": {
               "default": "",
               "title": "Teardown",
               "type": "string"
            }
         },
         "title": "BashTemplate",
         "type": "object"
      }
   },
   "additionalProperties": false,
   "required": [
      "script_template"
   ]
}

Config:

extra: str = forbid

Fields:

concurrency (int)
script_template (ai2_kit.core.script.BashTemplate)

field concurrency: int = 5#

field script_template: BashTemplate [Required]#

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])[source]#

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])[source]#

Bases: ICllExploreOutput

get_model_devi_dataset() → List[Artifact][source]#

output_dirs: List[Artifact]#

async ai2_kit.domain.anyware.anyware(input: AnywareInput, ctx: AnywareContext) → AnywareOutput[source]#

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)[source]#

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