ai2_kit.domain.deepmd module#

class ai2_kit.domain.deepmd.CllDeepmdContext(path_prefix: str, resource_manager: ai2_kit.core.resource_manager.ResourceManager, config: ai2_kit.domain.deepmd.CllDeepmdContextConfig)[source]#

Bases: BaseCllContext

config: CllDeepmdContextConfig#

pydantic model ai2_kit.domain.deepmd.CllDeepmdContextConfig[source]#

Bases: BaseModel

Show JSON schema

{
   "title": "CllDeepmdContextConfig",
   "type": "object",
   "properties": {
      "script_template": {
         "$ref": "#/$defs/BashTemplate"
      },
      "dp_cmd": {
         "default": "dp",
         "title": "Dp Cmd",
         "type": "string"
      },
      "concurrency": {
         "default": 5,
         "title": "Concurrency",
         "type": "integer"
      },
      "multi_gpus_per_job": {
         "default": false,
         "title": "Multi Gpus Per Job",
         "type": "boolean"
      }
   },
   "$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)
dp_cmd (str)
multi_gpus_per_job (bool)
script_template (ai2_kit.core.script.BashTemplate)

field concurrency: int = 5#

field dp_cmd: str = 'dp'#

field multi_gpus_per_job: bool = False#

field script_template: BashTemplate [Required]#

class ai2_kit.domain.deepmd.CllDeepmdInput(config: ai2_kit.domain.deepmd.CllDeepmdInputConfig, mode: Literal['default', 'dpff', 'fep-redox', 'fep-pka'], type_map: List[str], sel_type: List[int] | None, old_dataset: List[ai2_kit.core.artifact.Artifact], new_dataset: List[ai2_kit.core.artifact.Artifact], previous: List[ai2_kit.core.artifact.Artifact])[source]#

Bases: object

config: CllDeepmdInputConfig#

mode: Literal['default', 'dpff', 'fep-redox', 'fep-pka']#

new_dataset: List[Artifact]#

old_dataset: List[Artifact]#

previous: List[Artifact]#

sel_type: List[int] | None#

type_map: List[str]#

pydantic model ai2_kit.domain.deepmd.CllDeepmdInputConfig[source]#

Bases: BaseModel

Show JSON schema

{
   "title": "CllDeepmdInputConfig",
   "type": "object",
   "properties": {
      "train_dw": {
         "anyOf": [
            {
               "$ref": "#/$defs/DwTraningConfig"
            },
            {
               "type": "null"
            }
         ],
         "default": null
      },
      "model_num": {
         "default": 4,
         "title": "Model Num",
         "type": "integer"
      },
      "init_dataset": {
         "default": [],
         "items": {
            "type": "string"
         },
         "title": "Init Dataset",
         "type": "array"
      },
      "input_template": {
         "default": {},
         "title": "Input Template",
         "type": "object"
      },
      "compress_model": {
         "default": false,
         "title": "Compress Model",
         "type": "boolean"
      },
      "pretrained_model": {
         "anyOf": [
            {
               "type": "string"
            },
            {
               "type": "null"
            }
         ],
         "default": null,
         "title": "Pretrained Model"
      },
      "isolate_outliers": {
         "default": false,
         "title": "Isolate Outliers",
         "type": "boolean"
      },
      "outlier_f_cutoff": {
         "default": 10.0,
         "title": "Outlier F Cutoff",
         "type": "number"
      },
      "outlier_weight": {
         "default": 0.003,
         "title": "Outlier Weight",
         "type": "number"
      },
      "fixture_models": {
         "default": [],
         "items": {
            "type": "string"
         },
         "title": "Fixture Models",
         "type": "array"
      },
      "group_by_formula": {
         "default": false,
         "title": "Group By Formula",
         "type": "boolean"
      },
      "init_from_previous": {
         "default": false,
         "title": "Init From Previous",
         "type": "boolean"
      },
      "input_modifier_fn": {
         "anyOf": [
            {
               "type": "string"
            },
            {
               "type": "null"
            }
         ],
         "default": null,
         "title": "Input Modifier Fn"
      },
      "ignore_error": {
         "default": false,
         "title": "Ignore Error",
         "type": "boolean"
      },
      "dp_train_opts": {
         "default": "",
         "title": "Dp Train Opts",
         "type": "string"
      }
   },
   "$defs": {
      "DwTraningConfig": {
         "additionalProperties": false,
         "description": "Options for deep wannier model training",
         "properties": {
            "input_template": {
               "title": "Input Template",
               "type": "object"
            }
         },
         "required": [
            "input_template"
         ],
         "title": "DwTraningConfig",
         "type": "object"
      }
   },
   "additionalProperties": false
}

Config:

extra: str = forbid

Fields:

compress_model (bool)
dp_train_opts (str)
fixture_models (List[str])
group_by_formula (bool)
ignore_error (bool)
init_dataset (List[str])
init_from_previous (bool)
input_modifier_fn (str | None)
input_template (dict)
isolate_outliers (bool)
model_num (int)
outlier_f_cutoff (float)
outlier_weight (float)
pretrained_model (str | None)
train_dw (ai2_kit.domain.deepmd.CllDeepmdInputConfig.DwTraningConfig | None)

field compress_model: bool = False#: Whether to compress model after training.

field dp_train_opts: str = ''#: Extra options for dp train command

field fixture_models: List[str] = []#: Fixture models used to initialize training, support glob pattern. If this is not empty, then the whole training process will be skipped. This feature is useful for debugging, or explore more structures without training. The models should be on the remote executor. The name fixture is used as the concept of fixture in pytest.

field group_by_formula: bool = False#

Grouping dataset by formula If this is enabled, then the dataset will be grouped by formula. Otherwise, the dataset will be grouped by ancestor.

Set this to True when you have multiple structures with the same ancestor.

field ignore_error: bool = False#: Ignore non critical errors.

field init_dataset: List[str] = []#: Dataset used to initialize training.

field init_from_previous: bool = False#: Use the previous models to initialize the current training, which can speed up the training process.

field input_modifier_fn: str | None = None#

A python function to modify the input data. The function should take a input dict as input and return a dict.

def input_modifier_fn(input: dict) -> dict:: … return new_input

field input_template: dict = {}#: Deepmd input template.

field isolate_outliers: bool = False#: If isolate_outliers is enabled, then outlier data will be separated from training data.

field model_num: int = 4#: Total number of models to train.

field outlier_f_cutoff: float = 10.0#: The threshold of force magnitude to determine whether a data is outlier.

field outlier_weight: float = 0.003#: The weight of outlier data in training data.

field pretrained_model: str | None = None#: Pretrained model used to finetune.

field train_dw: DwTraningConfig | None = None#: Options for deep wannier model training, if None, then deep wannier model will not be trained.

pydantic model DwTraningConfig[source]#

Bases: BaseModel

Options for deep wannier model training

Show JSON schema

{
   "title": "DwTraningConfig",
   "description": "Options for deep wannier model training",
   "type": "object",
   "properties": {
      "input_template": {
         "title": "Input Template",
         "type": "object"
      }
   },
   "additionalProperties": false,
   "required": [
      "input_template"
   ]
}

Config:

extra: str = forbid

Fields:

input_template (dict)

field input_template: dict [Required]#: Deepmd input template. Ref: https://docs.deepmodeling.com/projects/deepmd/en/master/model/dplr.html

class ai2_kit.domain.deepmd.GenericDeepmdOutput(models: List[ai2_kit.core.artifact.Artifact], dataset: List[ai2_kit.core.artifact.Artifact])[source]#

Bases: ICllTrainOutput

dataset: List[Artifact]#

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

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

models: List[Artifact]#

async ai2_kit.domain.deepmd.cll_deepmd(input: CllDeepmdInput, ctx: CllDeepmdContext)[source]#

ai2_kit.domain.deepmd.make_deepmd_dataset(dataset_dir: str, outlier_dir: str, raw_data_collection: List[ArtifactDict], isolate_outliers: bool, outlier_f_cutoff: float, type_map: List[str], deepmd_input_template: dict, group_by_formula: bool, mode: str, sel_type: List[int] | None, ignore_error: bool)[source]#

ai2_kit.domain.deepmd.make_deepmd_input(input_template: dict, type_map: List[str], train_systems: List[str], outlier_systems: List[str], validation_systems: List[str], isolate_outliers: bool, outlier_weight: float)[source]#

ai2_kit.domain.deepmd.make_deepmd_task_dirs(input_template: dict, model_num: int, type_map: List[str], train_systems: List[str], outlier_systems: List[str], validation_systems: List[str], isolate_outliers: bool, outlier_weight: float, dw_input_template: dict | None, base_dir: str, input_modifier_fn: str | None)[source]#