BladeDISC DeePMD code

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/compress.py

+"""Compress a model, which including tabulating the embedding-net."""
+
+import os
+import json
+import logging
+from typing import Optional
+
+from deepmd.common import j_loader
+from deepmd.env import tf, GLOBAL_ENER_FLOAT_PRECISION
+from deepmd.utils.argcheck import normalize
+from deepmd.utils.compat import update_deepmd_input
+from deepmd.utils.errors import GraphTooLargeError, GraphWithoutTensorError
+from deepmd.utils.graph import get_tensor_by_name
+
+from .freeze import freeze
+from .train import train, get_rcut, get_min_nbor_dist
+from .transfer import transfer
+
+__all__ = ["compress"]
+
+log = logging.getLogger(__name__)
+
+
+def compress(
+    *,
+    input: str,
+    output: str,
+    extrapolate: int,
+    step: float,
+    frequency: str,
+    checkpoint_folder: str,
+    training_script: str,
+    mpi_log: str,
+    log_path: Optional[str],
+    log_level: int,
+    **kwargs
+):
+    """Compress model.
+
+    The table is composed of fifth-order polynomial coefficients and is assembled from
+    two sub-tables. The first table takes the step parameter as the domain's uniform step size,
+    while the second table takes 10 * step as it's uniform step size. The range of the
+    first table is automatically detected by the code, while the second table ranges
+    from the first table's upper boundary(upper) to the extrapolate(parameter) * upper.
+
+    Parameters
+    ----------
+    input : str
+        frozen model file to compress
+    output : str
+        compressed model filename
+    extrapolate : int
+        scale of model extrapolation
+    step : float
+        uniform step size of the tabulation's first table
+    frequency : str
+        frequency of tabulation overflow check
+    checkpoint_folder : str
+        trining checkpoint folder for freezing
+    training_script : str
+        training script of the input frozen model
+    mpi_log : str
+        mpi logging mode for training
+    log_path : Optional[str]
+        if speccified log will be written to this file
+    log_level : int
+        logging level
+    """
+    try:
+        t_jdata = get_tensor_by_name(input, 'train_attr/training_script')
+        t_min_nbor_dist = get_tensor_by_name(input, 'train_attr/min_nbor_dist')
+        jdata = json.loads(t_jdata)
+    except GraphWithoutTensorError as e:
+        if training_script == None:
+            raise RuntimeError(
+                "The input frozen model: %s has no training script or min_nbor_dist information, "
+                "which is not supported by the model compression interface. "
+                "Please consider using the --training-script command within the model compression interface to provide the training script of the input frozen model. "
+                "Note that the input training script must contain the correct path to the training data." % input
+            ) from e
+        elif not os.path.exists(training_script):
+            raise RuntimeError(
+                "The input training script %s (%s) does not exist! Please check the path of the training script. " % (input, os.path.abspath(input))
+            ) from e
+        else:
+            log.info("stage 0: compute the min_nbor_dist")
+            jdata = j_loader(training_script)
+            jdata = update_deepmd_input(jdata)
+            t_min_nbor_dist = get_min_nbor_dist(jdata, get_rcut(jdata))
+
+    _check_compress_type(input)
+
+    tf.constant(t_min_nbor_dist,
+        name = 'train_attr/min_nbor_dist',
+        dtype = GLOBAL_ENER_FLOAT_PRECISION)
+    jdata["model"]["compress"] = {}
+    jdata["model"]["compress"]["model_file"] = input
+    jdata["model"]["compress"]["min_nbor_dist"] = t_min_nbor_dist
+    jdata["model"]["compress"]["table_config"] = [
+        extrapolate,
+        step,
+        10 * step,
+        int(frequency),
+    ]
+    jdata["training"]["save_ckpt"] = os.path.join("model-compression", "model.ckpt")
+    jdata = update_deepmd_input(jdata)
+    jdata = normalize(jdata)
+
+    # check the descriptor info of the input file
+    # move to the specific Descriptor class
+
+    # stage 1: training or refining the model with tabulation
+    log.info("\n\n")
+    log.info("stage 1: compress the model")
+    control_file = "compress.json"
+    with open(control_file, "w") as fp:
+        json.dump(jdata, fp, indent=4)
+    try:
+        train(
+            INPUT=control_file,
+            init_model=None,
+            restart=None,
+            init_frz_model=None,
+            output=control_file,
+            mpi_log=mpi_log,
+            log_level=log_level,
+            log_path=log_path,
+            is_compress=True,
+        )
+    except GraphTooLargeError as e:
+        raise RuntimeError(
+            "The uniform step size of the tabulation's first table is %f, " 
+            "which is too small. This leads to a very large graph size, "
+            "exceeding protobuf's limitation (2 GB). You should try to "
+            "increase the step size." % step
+        ) from e
+
+    # reset the graph, otherwise the size limitation will be only 2 GB / 2 = 1 GB
+    tf.reset_default_graph()
+
+    # stage 2: freeze the model
+    log.info("\n\n")
+    log.info("stage 2: freeze the model")
+    try:
+        freeze(checkpoint_folder=checkpoint_folder, output=output, node_names=None)
+    except GraphTooLargeError as e:
+        raise RuntimeError(
+            "The uniform step size of the tabulation's first table is %f, " 
+            "which is too small. This leads to a very large graph size, "
+            "exceeding protobuf's limitation (2 GB). You should try to "
+            "increase the step size." % step
+        ) from e
+
+def _check_compress_type(model_file):
+    try:
+        t_model_type = bytes.decode(get_tensor_by_name(model_file, 'model_type'))
+    except GraphWithoutTensorError as e:
+        # Compatible with the upgraded model, which has no 'model_type' info
+        t_model_type = None
+    
+    if t_model_type == "compressed_model":
+        raise RuntimeError("The input frozen model %s has already been compressed! Please do not compress the model repeatedly. " % model_file)

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/config.py

+#!/usr/bin/env python3
+"""Quickly create a configuration file for smooth model."""
+
+import json
+import yaml
+from pathlib import Path
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+
+__all__ = ["config"]
+
+
+DEFAULT_DATA: Dict[str, Any] = {
+    "use_smooth": True,
+    "sel_a": [],
+    "rcut_smth": -1,
+    "rcut": -1,
+    "filter_neuron": [20, 40, 80],
+    "filter_resnet_dt": False,
+    "axis_neuron": 8,
+    "fitting_neuron": [240, 240, 240],
+    "fitting_resnet_dt": True,
+    "coord_norm": True,
+    "type_fitting_net": False,
+    "systems": [],
+    "set_prefix": "set",
+    "stop_batch": -1,
+    "batch_size": -1,
+    "start_lr": 0.001,
+    "decay_steps": -1,
+    "decay_rate": 0.95,
+    "start_pref_e": 0.02,
+    "limit_pref_e": 1,
+    "start_pref_f": 1000,
+    "limit_pref_f": 1,
+    "start_pref_v": 0,
+    "limit_pref_v": 0,
+    "seed": 1,
+    "disp_file": "lcurve.out",
+    "disp_freq": 1000,
+    "numb_test": 10,
+    "save_freq": 10000,
+    "save_ckpt": "model.ckpt",
+    "disp_training": True,
+    "time_training": True,
+}
+
+
+def valid_dir(path: Path):
+    """Check if directory is a valid deepmd system directory.
+
+    Parameters
+    ----------
+    path : Path
+        path to directory
+
+    Raises
+    ------
+    OSError
+        if `type.raw` is missing on dir or `box.npy` or `coord.npy` are missing in one
+        of the sets subdirs
+    """
+    if not (path / "type.raw").is_file():
+        raise OSError
+    for ii in path.glob("set.*"):
+        if not (ii / "box.npy").is_file():
+            raise OSError
+        if not (ii / "coord.npy").is_file():
+            raise OSError
+
+
+def load_systems(dirs: List[Path]) -> Tuple[List[np.ndarray], List[np.ndarray]]:
+    """Load systems to memory for disk.
+
+    Parameters
+    ----------
+    dirs : List[Path]
+        list of system directories paths
+
+    Returns
+    -------
+    Tuple[List[np.ndarray], List[np.ndarray]]
+        atoms types and structure cells formated as Nx9 array
+    """
+    all_type = []
+    all_box = []
+    for d in dirs:
+        sys_type = np.loadtxt(d / "type.raw", dtype=int)
+        sys_box = np.vstack([np.load(s / "box.npy") for s in d.glob("set.*")])
+        all_type.append(sys_type)
+        all_box.append(sys_box)
+    return all_type, all_box
+
+
+def get_system_names() -> List[Path]:
+    """Get system directory paths from stdin.
+
+    Returns
+    -------
+    List[Path]
+        list of system directories paths
+    """
+    dirs = input("Enter system path(s) (seperated by space, wild card supported): \n")
+    system_dirs = []
+    for dir_str in dirs.split():
+        found_dirs = Path.cwd().glob(dir_str)
+        for d in found_dirs:
+            valid_dir(d)
+            system_dirs.append(d)
+
+    return system_dirs
+
+
+def get_rcut() -> float:
+    """Get rcut from stdin from user.
+
+    Returns
+    -------
+    float
+        input rcut lenght converted to float
+
+    Raises
+    ------
+    ValueError
+        if rcut is smaller than 0.0
+    """
+    dv = 6.0
+    rcut_input = input(f"Enter rcut (default {dv:.1f} A): \n")
+    try:
+        rcut = float(rcut_input)
+    except ValueError as e:
+        print(f"invalid rcut: {e} setting to default: {dv:.1f}")
+        rcut = dv
+    if rcut <= 0:
+        raise ValueError("rcut should be > 0")
+    return rcut
+
+
+def get_batch_size_rule() -> int:
+    """Get minimal batch size from user from stdin.
+
+    Returns
+    -------
+    int
+        size of the batch
+
+    Raises
+    ------
+    ValueError
+        if batch size is <= 0
+    """
+    dv = 32
+    matom_input = input(
+        f"Enter the minimal number of atoms in a batch (default {dv:d}: \n"
+    )
+    try:
+        matom = int(matom_input)
+    except ValueError as e:
+        print(f"invalid batch size: {e} setting to default: {dv:d}")
+        matom = dv
+    if matom <= 0:
+        raise ValueError("the number should be > 0")
+    return matom
+
+
+def get_stop_batch() -> int:
+    """Get stop batch from user from stdin.
+
+    Returns
+    -------
+    int
+        size of the batch
+
+    Raises
+    ------
+    ValueError
+        if stop batch is <= 0
+    """
+    dv = 1000000
+    sb_input = input(f"Enter the stop batch (default {dv:d}): \n")
+    try:
+        sb = int(sb_input)
+    except ValueError as e:
+        print(f"invalid stop batch: {e} setting to default: {dv:d}")
+        sb = dv
+    if sb <= 0:
+        raise ValueError("the number should be > 0")
+    return sb
+
+
+def get_ntypes(all_type: List[np.ndarray]) -> int:
+    """Count number of unique elements.
+
+    Parameters
+    ----------
+    all_type : List[np.ndarray]
+        list with arrays specifying elements of structures
+
+    Returns
+    -------
+    int
+        number of unique elements
+    """
+    return len(np.unique(all_type))
+
+
+def get_max_density(
+    all_type: List[np.ndarray], all_box: List[np.ndarray]
+) -> np.ndarray:
+    """Compute maximum density in suppliedd cells.
+
+    Parameters
+    ----------
+    all_type : List[np.ndarray]
+        list with arrays specifying elements of structures
+    all_box : List[np.ndarray]
+        list with arrays specifying cells for all structures
+
+    Returns
+    -------
+    float
+        maximum atom density in all supplies structures for each element individually
+    """
+    ntypes = get_ntypes(all_type)
+    all_max = []
+    for tt, bb in zip(all_type, all_box):
+        vv = np.reshape(bb, [-1, 3, 3])
+        vv = np.linalg.det(vv)
+        min_v = np.min(vv)
+        type_count = []
+        for ii in range(ntypes):
+            type_count.append(sum(tt == ii))
+        max_den = type_count / min_v
+        all_max.append(max_den)
+    all_max = np.max(all_max, axis=0)
+    return all_max
+
+
+def suggest_sel(
+    all_type: List[np.ndarray],
+    all_box: List[np.ndarray],
+    rcut: float,
+    ratio: float = 1.5,
+) -> List[int]:
+    """Suggest selection parameter.
+
+    Parameters
+    ----------
+    all_type : List[np.ndarray]
+        list with arrays specifying elements of structures
+    all_box : List[np.ndarray]
+        list with arrays specifying cells for all structures
+    rcut : float
+        cutoff radius
+    ratio : float, optional
+        safety margin to add to estimated value, by default 1.5
+
+    Returns
+    -------
+    List[int]
+        [description]
+    """
+    max_den = get_max_density(all_type, all_box)
+    return [int(ii) for ii in max_den * 4.0 / 3.0 * np.pi * rcut ** 3 * ratio]
+
+
+def suggest_batch_size(all_type: List[np.ndarray], min_atom: int) -> List[int]:
+    """Get suggestion for batch size.
+
+    Parameters
+    ----------
+    all_type : List[np.ndarray]
+        list with arrays specifying elements of structures
+    min_atom : int
+        minimal number of atoms in batch
+
+    Returns
+    -------
+    List[int]
+        suggested batch sizes for each system
+    """
+    bs = []
+    for ii in all_type:
+        natoms = len(ii)
+        tbs = min_atom // natoms
+        if (min_atom // natoms) * natoms != min_atom:
+            tbs += 1
+        bs.append(tbs)
+    return bs
+
+
+def suggest_decay(stop_batch: int) -> Tuple[int, float]:
+    """Suggest number of decay steps and decay rate.
+
+    Parameters
+    ----------
+    stop_batch : int
+        stop batch number
+
+    Returns
+    -------
+    Tuple[int, float]
+        number of decay steps and decay rate
+    """
+    decay_steps = int(stop_batch // 200)
+    decay_rate = 0.95
+    return decay_steps, decay_rate
+
+
+def config(*, output: str, **kwargs):
+    """Auto config file generator.
+
+    Parameters
+    ----------
+    output: str
+        file to write config file
+
+    Raises
+    ------
+    RuntimeError
+        if user does not input any systems
+    ValueError
+        if output file is of wrong type
+    """
+    all_sys = get_system_names()
+    if len(all_sys) == 0:
+        raise RuntimeError("no system specified")
+    rcut = get_rcut()
+    matom = get_batch_size_rule()
+    stop_batch = get_stop_batch()
+
+    all_type, all_box = load_systems(all_sys)
+    sel = suggest_sel(all_type, all_box, rcut, ratio=1.5)
+    bs = suggest_batch_size(all_type, matom)
+    decay_steps, decay_rate = suggest_decay(stop_batch)
+
+    jdata = DEFAULT_DATA.copy()
+    jdata["systems"] = [str(ii) for ii in all_sys]
+    jdata["sel_a"] = sel
+    jdata["rcut"] = rcut
+    jdata["rcut_smth"] = rcut - 0.2
+    jdata["stop_batch"] = stop_batch
+    jdata["batch_size"] = bs
+    jdata["decay_steps"] = decay_steps
+    jdata["decay_rate"] = decay_rate
+
+    with open(output, "w") as fp:
+        if output.endswith("json"):
+            json.dump(jdata, fp, indent=4)
+        elif output.endswith(("yml", "yaml")):
+            yaml.safe_dump(jdata, fp, default_flow_style=False)
+        else:
+            raise ValueError("output file must be of type json or yaml")

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/convert.py

+from deepmd.utils.convert import convert_012_to_21, convert_10_to_21, convert_20_to_21, convert_13_to_21, convert_12_to_21 
+
+def convert(
+    *,
+    FROM: str,
+    input_model: str,
+    output_model: str,
+    **kwargs,
+):
+    if FROM == '0.12':
+        convert_012_to_21(input_model, output_model)
+    elif FROM == '1.0':
+        convert_10_to_21(input_model, output_model)
+    elif FROM in ['1.1', '1.2']:
+        # no difference between 1.1 and 1.2
+        convert_12_to_21(input_model, output_model)
+    elif FROM == '1.3':
+        convert_13_to_21(input_model, output_model)
+    elif FROM == '2.0':
+        convert_20_to_21(input_model, output_model)
+    else:
+        raise RuntimeError('unsupported model version ' + FROM)

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/doc.py

+"""Module that prints train input arguments docstrings."""
+
+from deepmd.utils.argcheck import gen_doc, gen_json
+
+__all__ = ["doc_train_input"]
+
+
+def doc_train_input(*, out_type: str = "rst", **kwargs):
+    """Print out trining input arguments to console."""
+    if out_type == "rst":
+        doc_str = gen_doc(make_anchor=True)
+    elif out_type == "json":
+        doc_str = gen_json()
+    else:
+        raise RuntimeError("Unsupported out type %s" % out_type)
+    print(doc_str)

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/freeze.py

+#!/usr/bin/env python3
+"""Script for freezing TF trained graph so it can be used with LAMMPS and i-PI.
+
+References
+----------
+https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc
+"""
+
+import logging
+import google.protobuf.message
+from deepmd.env import tf, FITTING_NET_PATTERN
+from deepmd.utils.errors import GraphTooLargeError
+from deepmd.utils.sess import run_sess
+from deepmd.utils.graph import get_pattern_nodes_from_graph_def
+from os.path import abspath
+
+# load grad of force module
+import deepmd.op
+
+from typing import List, Optional
+
+from deepmd.nvnmd.entrypoints.freeze import save_weight
+
+__all__ = ["freeze"]
+
+log = logging.getLogger(__name__)
+
+def _transfer_fitting_net_trainable_variables(sess, old_graph_def, raw_graph_def):
+    old_pattern = FITTING_NET_PATTERN
+    raw_pattern = FITTING_NET_PATTERN\
+        .replace('idt',    'idt+_\d+')\
+        .replace('bias',   'bias+_\d+')\
+        .replace('matrix', 'matrix+_\d+')
+    old_graph_nodes = get_pattern_nodes_from_graph_def(
+        old_graph_def, 
+        old_pattern
+    )
+    try :
+        raw_graph_def = tf.graph_util.convert_variables_to_constants(
+            sess,  # The session is used to retrieve the weights
+            raw_graph_def,  # The graph_def is used to retrieve the nodes
+            [n + '_1' for n in old_graph_nodes],  # The output node names are used to select the usefull nodes
+        )
+    except AssertionError:
+        # if there's no additional nodes
+        return old_graph_def
+
+    raw_graph_nodes = get_pattern_nodes_from_graph_def(
+        raw_graph_def, 
+        raw_pattern
+    )
+    for node in old_graph_def.node:
+        if node.name not in old_graph_nodes.keys():
+            continue
+        tensor = tf.make_ndarray(raw_graph_nodes[node.name + '_1'])
+        node.attr["value"].tensor.tensor_content = tensor.tostring()
+    return old_graph_def
+
+def _make_node_names(model_type: str, modifier_type: Optional[str] = None) -> List[str]:
+    """Get node names based on model type.
+
+    Parameters
+    ----------
+    model_type : str
+        str type of model
+    modifier_type : Optional[str], optional
+        modifier type if any, by default None
+
+    Returns
+    -------
+    List[str]
+        list with all node names to freeze
+
+    Raises
+    ------
+    RuntimeError
+        if unknown model type
+    """
+    nodes = [
+        "model_type",
+        "descrpt_attr/rcut",
+        "descrpt_attr/ntypes",
+        "model_attr/tmap",
+        "model_attr/model_type",
+        "model_attr/model_version",
+        "train_attr/min_nbor_dist",
+        "train_attr/training_script",
+    ]
+
+    if model_type == "ener":
+        nodes += [
+            "o_energy",
+            "o_force",
+            "o_virial",
+            "o_atom_energy",
+            "o_atom_virial",
+            "fitting_attr/dfparam",
+            "fitting_attr/daparam",
+        ]
+    elif model_type == "wfc":
+        nodes += [
+            "o_wfc",
+            "model_attr/sel_type",
+            "model_attr/output_dim",
+        ]
+    elif model_type == "dipole":
+        nodes += [
+            "o_dipole",
+            "o_global_dipole",
+            "o_force",
+            "o_virial",
+            "o_atom_virial",
+            "o_rmat",
+            "o_rmat_deriv",
+            "o_nlist",
+            "o_rij",
+            "descrpt_attr/sel",
+            "descrpt_attr/ndescrpt",
+            "model_attr/sel_type",
+            "model_attr/output_dim",
+        ]
+    elif model_type == "polar":
+        nodes += [
+            "o_polar",
+            "o_global_polar",
+            "o_force",
+            "o_virial",
+            "o_atom_virial",
+            "model_attr/sel_type",
+            "model_attr/output_dim",
+        ]
+    elif model_type == "global_polar":
+        nodes += [
+            "o_global_polar",
+            "model_attr/sel_type",
+            "model_attr/output_dim",
+        ]
+    else:
+        raise RuntimeError(f"unknow model type {model_type}")
+    if modifier_type == "dipole_charge":
+        nodes += [
+            "modifier_attr/type",
+            "modifier_attr/mdl_name",
+            "modifier_attr/mdl_charge_map",
+            "modifier_attr/sys_charge_map",
+            "modifier_attr/ewald_h",
+            "modifier_attr/ewald_beta",
+            "dipole_charge/model_type",
+            "dipole_charge/descrpt_attr/rcut",
+            "dipole_charge/descrpt_attr/ntypes",
+            "dipole_charge/model_attr/tmap",
+            "dipole_charge/model_attr/model_type",
+            "dipole_charge/model_attr/model_version",
+            "o_dm_force",
+            "dipole_charge/model_attr/sel_type",
+            "dipole_charge/o_dipole",
+            "dipole_charge/model_attr/output_dim",
+            "o_dm_virial",
+            "o_dm_av",
+        ]
+    return nodes
+
+
+def freeze(
+    *, checkpoint_folder: str, output: str, node_names: Optional[str] = None, nvnmd_weight: Optional[str] = None, **kwargs
+):
+    """Freeze the graph in supplied folder.
+
+    Parameters
+    ----------
+    checkpoint_folder : str
+        location of the folder with model
+    output : str
+        output file name
+    node_names : Optional[str], optional
+        names of nodes to output, by default None
+    """
+    # We retrieve our checkpoint fullpath
+    checkpoint = tf.train.get_checkpoint_state(checkpoint_folder)
+    input_checkpoint = checkpoint.model_checkpoint_path
+
+    # expand the output file to full path
+    output_graph = abspath(output)
+
+    # Before exporting our graph, we need to precise what is our output node
+    # This is how TF decides what part of the Graph he has to keep
+    # and what part it can dump
+    # NOTE: this variable is plural, because you can have multiple output nodes
+    # node_names = "energy_test,force_test,virial_test,t_rcut"
+
+    # We clear devices to allow TensorFlow to control
+    # on which device it will load operations
+    clear_devices = True
+
+    # We import the meta graph and retrieve a Saver
+    try:
+        # In case paralle training
+        import horovod.tensorflow as _
+    except ImportError:
+        pass
+    saver = tf.train.import_meta_graph(
+        f"{input_checkpoint}.meta", clear_devices=clear_devices
+    )
+
+    # We retrieve the protobuf graph definition
+    graph = tf.get_default_graph()
+    try:
+        input_graph_def = graph.as_graph_def()
+    except google.protobuf.message.DecodeError as e:
+        raise GraphTooLargeError(
+            "The graph size exceeds 2 GB, the hard limitation of protobuf."
+            " Then a DecodeError was raised by protobuf. You should "
+            "reduce the size of your model."
+        ) from e
+    nodes = [n.name for n in input_graph_def.node]
+
+    # We start a session and restore the graph weights
+    with tf.Session() as sess:
+        saver.restore(sess, input_checkpoint)
+        model_type = run_sess(sess, "model_attr/model_type:0", feed_dict={}).decode("utf-8")
+        if "modifier_attr/type" in nodes:
+            modifier_type = run_sess(sess, "modifier_attr/type:0", feed_dict={}).decode(
+                "utf-8"
+            )
+        else:
+            modifier_type = None
+        if node_names is None:
+            output_node_list = _make_node_names(model_type, modifier_type)
+            different_set = set(output_node_list) - set(nodes)
+            if different_set:
+                log.warning(
+                    "The following nodes are not in the graph: %s. "
+                    "Skip freezeing these nodes. You may be freezing "
+                    "a checkpoint generated by an old version." % different_set
+                )
+                # use intersection as output list
+                output_node_list = list(set(output_node_list) & set(nodes))
+        else:
+            output_node_list = node_names.split(",")
+        log.info(f"The following nodes will be frozen: {output_node_list}")
+
+        if nvnmd_weight is not None:
+            save_weight(sess, nvnmd_weight) # nvnmd
+
+        # We use a built-in TF helper to export variables to constants
+        output_graph_def = tf.graph_util.convert_variables_to_constants(
+            sess,  # The session is used to retrieve the weights
+            input_graph_def,  # The graph_def is used to retrieve the nodes
+            output_node_list,  # The output node names are used to select the usefull nodes
+        )
+
+        # If we need to transfer the fitting net variables
+        output_graph_def = _transfer_fitting_net_trainable_variables(
+            sess,
+            output_graph_def,
+            input_graph_def
+        )
+
+        # Finally we serialize and dump the output graph to the filesystem
+        with tf.gfile.GFile(output_graph, "wb") as f:
+            f.write(output_graph_def.SerializeToString())
+        log.info(f"{len(output_graph_def.node):d} ops in the final graph.")

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/main.py

+"""DeePMD-Kit entry point module."""
+
+import argparse
+import logging
+import textwrap
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from deepmd import __version__
+from deepmd.entrypoints import (
+    compress,
+    config,
+    doc_train_input,
+    freeze,
+    test,
+    train_dp,
+    transfer,
+    make_model_devi,
+    convert,
+    neighbor_stat,
+)
+from deepmd.loggers import set_log_handles
+
+from deepmd.nvnmd.entrypoints.train import train_nvnmd
+
+__all__ = ["main", "parse_args", "get_ll", "main_parser"]
+
+
+def get_ll(log_level: str) -> int:
+    """Convert string to python logging level.
+
+    Parameters
+    ----------
+    log_level : str
+        allowed input values are: DEBUG, INFO, WARNING, ERROR, 3, 2, 1, 0
+
+    Returns
+    -------
+    int
+        one of python logging module log levels - 10, 20, 30 or 40
+    """
+    if log_level.isdigit():
+        int_level = (4 - int(log_level)) * 10
+    else:
+        int_level = getattr(logging, log_level)
+
+    return int_level
+
+class RawTextArgumentDefaultsHelpFormatter(
+    argparse.RawTextHelpFormatter, argparse.ArgumentDefaultsHelpFormatter
+):
+    """This formatter is used to print multile-line help message with default value."""
+
+
+def main_parser() -> argparse.ArgumentParser:
+    """DeePMD-Kit commandline options argument parser.
+
+    Returns
+    -------
+    argparse.ArgumentParser
+        main parser of DeePMD-kit
+    """
+    parser = argparse.ArgumentParser(
+        description="DeePMD-kit: A deep learning package for many-body potential energy"
+        " representation and molecular dynamics",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    subparsers = parser.add_subparsers(title="Valid subcommands", dest="command")
+
+    # * logging options parser *********************************************************
+    # with use of the parent argument this options will be added to every parser
+    parser_log = argparse.ArgumentParser(
+        add_help=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser_log.add_argument(
+        "-v",
+        "--log-level",
+        choices=["DEBUG", "3", "INFO", "2", "WARNING", "1", "ERROR", "0"],
+        default="INFO",
+        help="set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO "
+        "and 3=DEBUG",
+    )
+    parser_log.add_argument(
+        "-l",
+        "--log-path",
+        type=str,
+        default=None,
+        help="set log file to log messages to disk, if not specified, the logs will "
+        "only be output to console",
+    )
+    # * mpi logging parser *************************************************************
+    parser_mpi_log = argparse.ArgumentParser(
+        add_help=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser_mpi_log.add_argument(
+        "-m",
+        "--mpi-log",
+        type=str,
+        default="master",
+        choices=("master", "collect", "workers"),
+        help="Set the manner of logging when running with MPI. 'master' logs only on "
+        "main process, 'collect' broadcasts logs from workers to master and 'workers' "
+        "means each process will output its own log",
+    )
+
+    # * config script ******************************************************************
+    parser_cfig = subparsers.add_parser(
+        "config",
+        parents=[parser_log],
+        help="fast configuration of parameter file for smooth model",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser_cfig.add_argument(
+        "-o", "--output", type=str, default="input.json", help="the output json file"
+    )
+
+    # * transfer script ****************************************************************
+    parser_transfer = subparsers.add_parser(
+        "transfer", parents=[parser_log], help="pass parameters to another model"
+    )
+    parser_transfer.add_argument(
+        "-r",
+        "--raw-model",
+        default="raw_frozen_model.pb",
+        type=str,
+        help="the model receiving parameters",
+    )
+    parser_transfer.add_argument(
+        "-O",
+        "--old-model",
+        default="old_frozen_model.pb",
+        type=str,
+        help="the model providing parameters",
+    )
+    parser_transfer.add_argument(
+        "-o",
+        "--output",
+        default="frozen_model.pb",
+        type=str,
+        help="the model after passing parameters",
+    )
+
+    # * config parser ******************************************************************
+    parser_train = subparsers.add_parser(
+        "train",
+        parents=[parser_log, parser_mpi_log],
+        help="train a model",
+        formatter_class=RawTextArgumentDefaultsHelpFormatter,
+        epilog=textwrap.dedent("""\
+        examples:
+            dp train input.json
+            dp train input.json --restart model.ckpt
+            dp train input.json --init-model model.ckpt
+        """),
+    )
+    parser_train.add_argument(
+        "INPUT", help="the input parameter file in json or yaml format"
+    )
+    parser_train.add_argument(
+        "-i",
+        "--init-model",
+        type=str,
+        default=None,
+        help="Initialize the model by the provided checkpoint.",
+    )
+    parser_train.add_argument(
+        "-r",
+        "--restart",
+        type=str,
+        default=None,
+        help="Restart the training from the provided checkpoint.",
+    )
+    parser_train.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        default="out.json",
+        help="The output file of the parameters used in training.",
+    )
+    parser_train.add_argument(
+        "-f",
+        "--init-frz-model",
+        type=str,
+        default=None,
+        help="Initialize the training from the frozen model.",
+    )
+    parser_train.add_argument(
+        "--skip-neighbor-stat",
+        action="store_true",
+        help="Skip calculating neighbor statistics. Sel checking, automatic sel, and model compression will be disabled.",
+    )
+
+    # * freeze script ******************************************************************
+    parser_frz = subparsers.add_parser(
+        "freeze",
+        parents=[parser_log],
+        help="freeze the model",
+        formatter_class=RawTextArgumentDefaultsHelpFormatter,
+        epilog=textwrap.dedent("""\
+        examples:
+            dp freeze
+            dp freeze -o graph.pb
+        """),
+    )
+    parser_frz.add_argument(
+        "-c",
+        "--checkpoint-folder",
+        type=str,
+        default=".",
+        help="path to checkpoint folder",
+    )
+    parser_frz.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        default="frozen_model.pb",
+        help="name of graph, will output to the checkpoint folder",
+    )
+    parser_frz.add_argument(
+        "-n",
+        "--node-names",
+        type=str,
+        default=None,
+        help="the frozen nodes, if not set, determined from the model type",
+    )
+    parser_frz.add_argument(
+        "-w",
+        "--nvnmd-weight",
+        type=str,
+        default=None,
+        help="the name of weight file (.npy), if set, save the model's weight into the file",
+    )
+
+    # * test script ********************************************************************
+    parser_tst = subparsers.add_parser(
+        "test",
+        parents=[parser_log],
+        help="test the model",
+        formatter_class=RawTextArgumentDefaultsHelpFormatter,
+        epilog=textwrap.dedent("""\
+        examples:
+            dp test -m graph.pb -s /path/to/system -n 30
+        """),
+    )
+    parser_tst.add_argument(
+        "-m",
+        "--model",
+        default="frozen_model.pb",
+        type=str,
+        help="Frozen model file to import",
+    )
+    parser_tst.add_argument(
+        "-s",
+        "--system",
+        default=".",
+        type=str,
+        help="The system dir. Recursively detect systems in this directory",
+    )
+    parser_tst.add_argument(
+        "-S", "--set-prefix", default="set", type=str, help="The set prefix"
+    )
+    parser_tst.add_argument(
+        "-n", "--numb-test", default=100, type=int, help="The number of data for test"
+    )
+    parser_tst.add_argument(
+        "-r", "--rand-seed", type=int, default=None, help="The random seed"
+    )
+    parser_tst.add_argument(
+        "--shuffle-test", action="store_true", default=False, help="Shuffle test data"
+    )
+    parser_tst.add_argument(
+        "-d",
+        "--detail-file",
+        type=str,
+        default=None,
+        help="File where details of energy force and virial accuracy will be written",
+    )
+    parser_tst.add_argument(
+        "-a",
+        "--atomic",
+        action="store_true",
+        default=False,
+        help="Test the accuracy of atomic label, i.e. energy / tensor (dipole, polar)",
+    )
+
+    # * compress model *****************************************************************
+    # Compress a model, which including tabulating the embedding-net.
+    # The table is composed of fifth-order polynomial coefficients and is assembled
+    # from two sub-tables. The first table takes the step(parameter) as it's uniform
+    # step, while the second table takes 10 * step as it\s uniform step
+    #  The range of the first table is automatically detected by deepmd-kit, while the
+    # second table ranges from the first table's upper boundary(upper) to the
+    # extrapolate(parameter) * upper.
+    parser_compress = subparsers.add_parser(
+        "compress",
+        parents=[parser_log, parser_mpi_log],
+        help="compress a model",
+        formatter_class=RawTextArgumentDefaultsHelpFormatter,
+        epilog=textwrap.dedent("""\
+        examples:
+            dp compress
+            dp compress -i graph.pb -o compressed.pb
+        """),
+    )
+    parser_compress.add_argument(
+        "-i",
+        "--input",
+        default="frozen_model.pb",
+        type=str,
+        help="The original frozen model, which will be compressed by the code",
+    )
+    parser_compress.add_argument(
+        "-o",
+        "--output",
+        default="frozen_model_compressed.pb",
+        type=str,
+        help="The compressed model",
+    )
+    parser_compress.add_argument(
+        "-s",
+        "--step",
+        default=0.01,
+        type=float,
+        help="Model compression uses fifth-order polynomials to interpolate the embedding-net. " 
+        "It introduces two tables with different step size to store the parameters of the polynomials. "
+        "The first table covers the range of the training data, while the second table is an extrapolation of the training data. "
+        "The domain of each table is uniformly divided by a given step size. "
+        "And the step(parameter) denotes the step size of the first table and the second table will "
+        "use 10 * step as it's step size to save the memory. "
+        "Usually the value ranges from 0.1 to 0.001. " 
+        "Smaller step means higher accuracy and bigger model size",
+    )
+    parser_compress.add_argument(
+        "-e",
+        "--extrapolate",
+        default=5,
+        type=int,
+        help="The domain range of the first table is automatically detected by the code: [d_low, d_up]. "
+        "While the second table ranges from the first table's upper boundary(d_up) to the extrapolate(parameter) * d_up: [d_up, extrapolate * d_up]",
+    )
+    parser_compress.add_argument(
+        "-f",
+        "--frequency",
+        default=-1,
+        type=int,
+        help="The frequency of tabulation overflow check(Whether the input environment "
+        "matrix overflow the first or second table range). "
+        "By default do not check the overflow",
+    )
+    parser_compress.add_argument(
+        "-c",
+        "--checkpoint-folder",
+        type=str,
+        default="model-compression",
+        help="path to checkpoint folder",
+    )
+    parser_compress.add_argument(
+        "-t",
+        "--training-script",
+        type=str,
+        default=None,
+        help="The training script of the input frozen model",
+    )
+
+    # * print docs script **************************************************************
+    parsers_doc = subparsers.add_parser(
+        "doc-train-input",
+        parents=[parser_log],
+        help="print the documentation (in rst format) of input training parameters.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parsers_doc.add_argument(
+        "--out-type", 
+        default="rst", 
+        type=str, 
+        help="The output type"
+    )
+
+    # * make model deviation ***********************************************************
+    parser_model_devi = subparsers.add_parser(
+        "model-devi",
+        parents=[parser_log],
+        help="calculate model deviation",
+        formatter_class=RawTextArgumentDefaultsHelpFormatter,
+        epilog=textwrap.dedent("""\
+        examples:
+            dp model-devi -m graph.000.pb graph.001.pb graph.002.pb graph.003.pb -s ./data -o model_devi.out
+        """),
+    )
+    parser_model_devi.add_argument(
+        "-m",
+        "--models",
+        default=["graph.000.pb", "graph.001.pb", "graph.002.pb", "graph.003.pb"],
+        nargs="+",
+        type=str,
+        help="Frozen models file to import",
+    )
+    parser_model_devi.add_argument(
+        "-s",
+        "--system",
+        default=".",
+        type=str,
+        help="The system directory. Recursively detect systems in this directory.",
+    )
+    parser_model_devi.add_argument(
+        "-S", "--set-prefix", default="set", type=str, help="The set prefix"
+    )
+    parser_model_devi.add_argument(
+        "-o",
+        "--output", 
+        default="model_devi.out", 
+        type=str, 
+        help="The output file for results of model deviation"
+    )
+    parser_model_devi.add_argument(
+        "-f",
+        "--frequency",
+        default=1,
+        type=int,
+        help="The trajectory frequency of the system"
+    )
+
+    # * convert models
+    parser_transform = subparsers.add_parser(
+        'convert-from',
+        parents=[parser_log],
+        help='convert lower model version to supported version',
+        formatter_class=RawTextArgumentDefaultsHelpFormatter,
+        epilog=textwrap.dedent("""\
+        examples:
+            dp convert-from 1.0 -i graph.pb -o graph_new.pb
+        """),
+    )
+    parser_transform.add_argument(
+        'FROM',
+        type = str,
+        choices = ['0.12', '1.0', '1.1', '1.2', '1.3', '2.0'],
+        help="The original model compatibility",
+    )
+    parser_transform.add_argument(
+        '-i',
+        "--input-model",
+        default = "frozen_model.pb",
+        type=str, 
+		help = "the input model",
+    )
+    parser_transform.add_argument(
+        "-o",
+        "--output-model",
+        default = "convert_out.pb",
+        type=str, 
+		help='the output model',
+    )
+
+    # neighbor_stat
+    parser_neighbor_stat = subparsers.add_parser(
+        'neighbor-stat',
+        parents=[parser_log],
+        help='Calculate neighbor statistics',
+        formatter_class=RawTextArgumentDefaultsHelpFormatter,
+        epilog=textwrap.dedent("""\
+        examples:
+            dp neighbor-stat -s data -r 6.0 -t O H
+        """),
+    )
+    parser_neighbor_stat.add_argument(
+        "-s",
+        "--system",
+        default=".",
+        type=str,
+        help="The system dir. Recursively detect systems in this directory",
+    )
+    parser_neighbor_stat.add_argument(
+        "-r",
+        "--rcut",
+        type=float,
+        required=True,
+        help="cutoff radius",
+    )
+    parser_neighbor_stat.add_argument(
+        "-t",
+        "--type-map",
+        type=str,
+        nargs='+',
+        required=True,
+        help="type map",
+    )
+    parser_neighbor_stat.add_argument(
+        "--one-type",
+        action="store_true",
+        default=False,
+        help="treat all types as a single type. Used with se_atten descriptor.",
+    )
+
+    # --version
+    parser.add_argument('--version', action='version', version='DeePMD-kit v%s' % __version__)
+
+    # * train nvnmd script ******************************************************************
+    parser_train_nvnmd = subparsers.add_parser(
+        "train-nvnmd",
+        parents=[parser_log],
+        help="train nvnmd model",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser_train_nvnmd.add_argument(
+        "INPUT", help="the input parameter file in json format"
+    )
+    parser_train_nvnmd.add_argument(
+        "-s",
+        "--step",
+        default="s1",
+        type=str,
+        choices=['s1', 's2'],
+        help="steps to train model of NVNMD: s1 (train CNN), s2 (train QNN)"
+    )
+    return parser
+
+
+def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace:
+    """Parse arguments and convert argument strings to objects.
+
+    Parameters
+    ----------
+    args: List[str]
+        list of command line arguments, main purpose is testing default option None
+        takes arguments from sys.argv
+
+    Returns
+    -------
+    argparse.Namespace
+        the populated namespace
+    """
+    parser = main_parser()
+    parsed_args = parser.parse_args(args=args)
+    if parsed_args.command is None:
+        parser.print_help()
+    else:
+        parsed_args.log_level = get_ll(parsed_args.log_level)
+
+    return parsed_args
+
+
+def main():
+    """DeePMD-Kit entry point.
+
+    Raises
+    ------
+    RuntimeError
+        if no command was input
+    """
+    args = parse_args()
+
+    # do not set log handles for None, it is useless
+    # log handles for train will be set separatelly
+    # when the use of MPI will be determined in `RunOptions`
+    if args.command not in (None, "train"):
+        set_log_handles(args.log_level, Path(args.log_path) if args.log_path else None)
+
+    dict_args = vars(args)
+
+    if args.command == "train":
+        train_dp(**dict_args)
+    elif args.command == "freeze":
+        freeze(**dict_args)
+    elif args.command == "config":
+        config(**dict_args)
+    elif args.command == "test":
+        test(**dict_args)
+    elif args.command == "transfer":
+        transfer(**dict_args)
+    elif args.command == "compress":
+        compress(**dict_args)
+    elif args.command == "doc-train-input":
+        doc_train_input(**dict_args)
+    elif args.command == "model-devi":
+        make_model_devi(**dict_args)
+    elif args.command == "convert-from":
+        convert(**dict_args)
+    elif args.command == "neighbor-stat":
+        neighbor_stat(**dict_args)
+    elif args.command == "train-nvnmd":  # nvnmd
+        train_nvnmd(**dict_args)
+    elif args.command is None:
+        pass
+    else:
+        raise RuntimeError(f"unknown command {args.command}")

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/neighbor_stat.py

+import logging
+from typing import List
+
+from deepmd.common import expand_sys_str
+from deepmd.utils.data_system import DeepmdDataSystem
+from deepmd.utils.neighbor_stat import NeighborStat
+
+log = logging.getLogger(__name__)
+
+def neighbor_stat(
+    *,
+    system: str,
+    rcut: float,
+    type_map: List[str],
+    one_type: bool = False,
+    **kwargs,
+):
+    """Calculate neighbor statistics.
+
+    Parameters
+    ----------
+    system : str
+        system to stat
+    rcut : float
+        cutoff radius
+    type_map : list[str]
+        type map
+    one_type : bool, optional, default=False
+        treat all types as a single type
+
+    Examples
+    --------
+    >>> neighbor_stat(system='.', rcut=6., type_map=["C", "H", "O", "N", "P", "S", "Mg", "Na", "HW", "OW", "mNa", "mCl", "mC", "mH", "mMg", "mN", "mO", "mP"])
+    min_nbor_dist: 0.6599510670195264
+    max_nbor_size: [23, 26, 19, 16, 2, 2, 1, 1, 72, 37, 5, 0, 31, 29, 1, 21, 20, 5]
+    """ 
+    all_sys = expand_sys_str(system)
+    if not len(all_sys):
+        raise RuntimeError("Did not find valid system")
+    data = DeepmdDataSystem(
+        systems=all_sys,
+        batch_size=1,
+        test_size=1,
+        rcut=rcut,
+        type_map=type_map,
+    )
+    data.get_batch()
+    nei = NeighborStat(data.get_ntypes(), rcut, one_type=one_type)
+    min_nbor_dist, max_nbor_size = nei.get_stat(data)
+    log.info("min_nbor_dist: %f" % min_nbor_dist)
+    log.info("max_nbor_size: %s" % str(max_nbor_size))
+    return min_nbor_dist, max_nbor_size

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/test.py

+"""Test trained DeePMD model."""
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, List, Dict, Optional, Tuple
+
+import numpy as np
+from deepmd import DeepPotential
+from deepmd.common import expand_sys_str
+from deepmd.utils import random as dp_random
+from deepmd.utils.data import DeepmdData
+from deepmd.utils.weight_avg import weighted_average
+
+if TYPE_CHECKING:
+    from deepmd.infer import DeepDipole, DeepPolar, DeepPot, DeepWFC
+    from deepmd.infer.deep_tensor import DeepTensor
+
+__all__ = ["test"]
+
+log = logging.getLogger(__name__)
+
+
+def test(
+    *,
+    model: str,
+    system: str,
+    set_prefix: str,
+    numb_test: int,
+    rand_seed: Optional[int],
+    shuffle_test: bool,
+    detail_file: str,
+    atomic: bool,
+    **kwargs,
+):
+    """Test model predictions.
+
+    Parameters
+    ----------
+    model : str
+        path where model is stored
+    system : str
+        system directory
+    set_prefix : str
+        string prefix of set
+    numb_test : int
+        munber of tests to do
+    rand_seed : Optional[int]
+        seed for random generator
+    shuffle_test : bool
+        whether to shuffle tests
+    detail_file : Optional[str]
+        file where test details will be output
+    atomic : bool
+        whether per atom quantities should be computed
+
+    Raises
+    ------
+    RuntimeError
+        if no valid system was found
+    """
+    all_sys = expand_sys_str(system)
+    if len(all_sys) == 0:
+        raise RuntimeError("Did not find valid system")
+    err_coll = []
+    siz_coll = []
+
+    # init random seed
+    if rand_seed is not None:
+        dp_random.seed(rand_seed % (2 ** 32))
+
+    # init model
+    dp = DeepPotential(model)
+
+    for cc, system in enumerate(all_sys):
+        log.info("# ---------------output of dp test--------------- ")
+        log.info(f"# testing system : {system}")
+
+        # create data class
+        tmap = dp.get_type_map() if dp.model_type == "ener" else None
+        data = DeepmdData(system, set_prefix, shuffle_test=shuffle_test, type_map=tmap)
+
+        if dp.model_type == "ener":
+            err = test_ener(
+                dp,
+                data,
+                system,
+                numb_test,
+                detail_file,
+                atomic,
+                append_detail=(cc != 0),
+            )
+        elif dp.model_type == "dipole":
+            err = test_dipole(dp, data, numb_test, detail_file, atomic)
+        elif dp.model_type == "polar":
+            err = test_polar(dp, data, numb_test, detail_file, atomic=atomic)
+        elif dp.model_type == "global_polar":   # should not appear in this new version
+            log.warning("Global polar model is not currently supported. Please directly use the polar mode and change loss parameters.")
+            err = test_polar(dp, data, numb_test, detail_file, atomic=False)    # YWolfeee: downward compatibility
+        log.info("# ----------------------------------------------- ")
+        err_coll.append(err)
+
+    avg_err = weighted_average(err_coll)
+
+    if len(all_sys) != len(err_coll):
+        log.warning("Not all systems are tested! Check if the systems are valid")
+
+    if len(all_sys) > 1:
+        log.info("# ----------weighted average of errors----------- ")
+        log.info(f"# number of systems : {len(all_sys)}")
+        if dp.model_type == "ener":
+            print_ener_sys_avg(avg_err)
+        elif dp.model_type == "dipole":
+            print_dipole_sys_avg(avg_err)
+        elif dp.model_type == "polar":
+            print_polar_sys_avg(avg_err)
+        elif dp.model_type == "global_polar":
+            print_polar_sys_avg(avg_err)
+        elif dp.model_type == "wfc":
+            print_wfc_sys_avg(avg_err)
+        log.info("# ----------------------------------------------- ")
+
+
+def rmse(diff: np.ndarray) -> np.ndarray:
+    """Calculate average root mean square error.
+
+    Parameters
+    ----------
+    diff: np.ndarray
+        difference
+
+    Returns
+    -------
+    np.ndarray
+        array with normalized difference
+    """
+    return np.sqrt(np.average(diff * diff))
+
+
+def save_txt_file(
+    fname: Path, data: np.ndarray, header: str = "", append: bool = False
+):
+    """Save numpy array to test file.
+
+    Parameters
+    ----------
+    fname : str
+        filename
+    data : np.ndarray
+        data to save to disk
+    header : str, optional
+        header string to use in file, by default ""
+    append : bool, optional
+        if true file will be appended insted of overwriting, by default False
+    """
+    flags = "ab" if append else "w"
+    with fname.open(flags) as fp:
+        np.savetxt(fp, data, header=header)
+
+
+def test_ener(
+    dp: "DeepPot",
+    data: DeepmdData,
+    system: str,
+    numb_test: int,
+    detail_file: Optional[str],
+    has_atom_ener: bool,
+    append_detail: bool = False,
+) -> Tuple[List[np.ndarray], List[int]]:
+    """Test energy type model.
+
+    Parameters
+    ----------
+    dp : DeepPot
+        instance of deep potential
+    data: DeepmdData
+        data container object
+    system : str
+        system directory
+    numb_test : int
+        munber of tests to do
+    detail_file : Optional[str]
+        file where test details will be output
+    has_atom_ener : bool
+        whether per atom quantities should be computed
+    append_detail : bool, optional
+        if true append output detail file, by default False
+
+    Returns
+    -------
+    Tuple[List[np.ndarray], List[int]]
+        arrays with results and their shapes
+    """
+    data.add("energy", 1, atomic=False, must=False, high_prec=True)
+    data.add("force", 3, atomic=True, must=False, high_prec=False)
+    data.add("virial", 9, atomic=False, must=False, high_prec=False)
+    if dp.has_efield:
+        data.add("efield", 3, atomic=True, must=True, high_prec=False)
+    if has_atom_ener:
+        data.add("atom_ener", 1, atomic=True, must=True, high_prec=False)
+    if dp.get_dim_fparam() > 0:
+        data.add(
+            "fparam", dp.get_dim_fparam(), atomic=False, must=True, high_prec=False
+        )
+    if dp.get_dim_aparam() > 0:
+        data.add("aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False)
+
+    test_data = data.get_test()
+    natoms = len(test_data["type"][0])
+    nframes = test_data["box"].shape[0]
+    numb_test = min(nframes, numb_test)
+
+    coord = test_data["coord"][:numb_test].reshape([numb_test, -1])
+    box = test_data["box"][:numb_test]
+    if dp.has_efield:
+        efield = test_data["efield"][:numb_test].reshape([numb_test, -1])
+    else:
+        efield = None
+    if not data.pbc:
+        box = None
+    atype = test_data["type"][0]
+    if dp.get_dim_fparam() > 0:
+        fparam = test_data["fparam"][:numb_test]
+    else:
+        fparam = None
+    if dp.get_dim_aparam() > 0:
+        aparam = test_data["aparam"][:numb_test]
+    else:
+        aparam = None
+
+    ret = dp.eval(
+        coord,
+        box,
+        atype,
+        fparam=fparam,
+        aparam=aparam,
+        atomic=has_atom_ener,
+        efield=efield,
+    )
+    energy = ret[0]
+    force = ret[1]
+    virial = ret[2]
+    energy = energy.reshape([numb_test, 1])
+    force = force.reshape([numb_test, -1])
+    virial = virial.reshape([numb_test, 9])
+    if has_atom_ener:
+        ae = ret[3]
+        av = ret[4]
+        ae = ae.reshape([numb_test, -1])
+        av = av.reshape([numb_test, -1])
+
+    rmse_e = rmse(energy - test_data["energy"][:numb_test].reshape([-1, 1]))
+    rmse_f = rmse(force - test_data["force"][:numb_test])
+    rmse_v = rmse(virial - test_data["virial"][:numb_test])
+    rmse_ea = rmse_e / natoms
+    rmse_va = rmse_v / natoms
+    if has_atom_ener:
+        rmse_ae = rmse(
+            test_data["atom_ener"][:numb_test].reshape([-1]) - ae.reshape([-1])
+        )
+
+    # print ("# energies: %s" % energy)
+    log.info(f"# number of test data : {numb_test:d} ")
+    log.info(f"Energy RMSE        : {rmse_e:e} eV")
+    log.info(f"Energy RMSE/Natoms : {rmse_ea:e} eV")
+    log.info(f"Force  RMSE        : {rmse_f:e} eV/A")
+    if data.pbc:
+        log.info(f"Virial RMSE        : {rmse_v:e} eV")
+        log.info(f"Virial RMSE/Natoms : {rmse_va:e} eV")
+    if has_atom_ener:
+        log.info(f"Atomic ener RMSE   : {rmse_ae:e} eV")
+
+    if detail_file is not None:
+        detail_path = Path(detail_file)
+
+        pe = np.concatenate(
+            (
+                np.reshape(test_data["energy"][:numb_test], [-1, 1]),
+                np.reshape(energy, [-1, 1]),
+            ),
+            axis=1,
+        )
+        save_txt_file(
+            detail_path.with_suffix(".e.out"),
+            pe,
+            header="%s: data_e pred_e" % system,
+            append=append_detail,
+        )
+        pf = np.concatenate(
+            (
+                np.reshape(test_data["force"][:numb_test], [-1, 3]),
+                np.reshape(force, [-1, 3]),
+            ),
+            axis=1,
+        )
+        save_txt_file(
+            detail_path.with_suffix(".f.out"),
+            pf,
+            header="%s: data_fx data_fy data_fz pred_fx pred_fy pred_fz" % system,
+            append=append_detail,
+        )
+        pv = np.concatenate(
+            (
+                np.reshape(test_data["virial"][:numb_test], [-1, 9]),
+                np.reshape(virial, [-1, 9]),
+            ),
+            axis=1,
+        )
+        save_txt_file(
+            detail_path.with_suffix(".v.out"),
+            pv,
+            header=f"{system}: data_vxx data_vxy data_vxz data_vyx data_vyy "
+            "data_vyz data_vzx data_vzy data_vzz pred_vxx pred_vxy pred_vxz pred_vyx "
+            "pred_vyy pred_vyz pred_vzx pred_vzy pred_vzz",
+            append=append_detail,
+        )
+    return {
+        "rmse_ea" : (rmse_ea, energy.size),
+        "rmse_f" : (rmse_f, force.size),
+        "rmse_va" : (rmse_va, virial.size),
+    }
+
+
+def print_ener_sys_avg(avg: Dict[str,float]):
+    """Print errors summary for energy type potential.
+
+    Parameters
+    ----------
+    avg : np.ndarray
+        array with summaries
+    """
+    log.info(f"Energy RMSE/Natoms : {avg['rmse_ea']:e} eV")
+    log.info(f"Force  RMSE        : {avg['rmse_f']:e} eV/A")
+    log.info(f"Virial RMSE/Natoms : {avg['rmse_va']:e} eV")
+
+
+def run_test(dp: "DeepTensor", test_data: dict, numb_test: int):
+    """Run tests.
+
+    Parameters
+    ----------
+    dp : DeepTensor
+        instance of deep potential
+    test_data : dict
+        dictionary with test data
+    numb_test : int
+        munber of tests to do
+
+    Returns
+    -------
+    [type]
+        [description]
+    """
+    nframes = test_data["box"].shape[0]
+    numb_test = min(nframes, numb_test)
+
+    coord = test_data["coord"][:numb_test].reshape([numb_test, -1])
+    box = test_data["box"][:numb_test]
+    atype = test_data["type"][0]
+    prediction = dp.eval(coord, box, atype)
+
+    return prediction.reshape([numb_test, -1]), numb_test, atype
+
+
+def test_wfc(
+    dp: "DeepWFC",
+    data: DeepmdData,
+    numb_test: int,
+    detail_file: Optional[str],
+) -> Tuple[List[np.ndarray], List[int]]:
+    """Test energy type model.
+
+    Parameters
+    ----------
+    dp : DeepPot
+        instance of deep potential
+    data: DeepmdData
+        data container object
+    numb_test : int
+        munber of tests to do
+    detail_file : Optional[str]
+        file where test details will be output
+
+    Returns
+    -------
+    Tuple[List[np.ndarray], List[int]]
+        arrays with results and their shapes
+    """
+    data.add(
+        "wfc", 12, atomic=True, must=True, high_prec=False, type_sel=dp.get_sel_type()
+    )
+    test_data = data.get_test()
+    wfc, numb_test, _ = run_test(dp, test_data, numb_test)
+    rmse_f = rmse(wfc - test_data["wfc"][:numb_test])
+
+    log.info("# number of test data : {numb_test:d} ")
+    log.info("WFC  RMSE : {rmse_f:e} eV/A")
+
+    if detail_file is not None:
+        detail_path = Path(detail_file)
+        pe = np.concatenate(
+            (
+                np.reshape(test_data["wfc"][:numb_test], [-1, 12]),
+                np.reshape(wfc, [-1, 12]),
+            ),
+            axis=1,
+        )
+        np.savetxt(
+            detail_path.with_suffix(".out"),
+            pe,
+            header="ref_wfc(12 dofs)   predicted_wfc(12 dofs)",
+        )
+    return {
+        'rmse' : (rmse_f, wfc.size)
+    }
+
+
+def print_wfc_sys_avg(avg):
+    """Print errors summary for wfc type potential.
+
+    Parameters
+    ----------
+    avg : np.ndarray
+        array with summaries
+    """
+    log.info(f"WFC  RMSE : {avg['rmse']:e} eV/A")
+
+
+def test_polar(
+    dp: "DeepPolar",
+    data: DeepmdData,
+    numb_test: int,
+    detail_file: Optional[str],
+    *,
+    atomic: bool,
+) -> Tuple[List[np.ndarray], List[int]]:
+    """Test energy type model.
+
+    Parameters
+    ----------
+    dp : DeepPot
+        instance of deep potential
+    data: DeepmdData
+        data container object
+    numb_test : int
+        munber of tests to do
+    detail_file : Optional[str]
+        file where test details will be output
+    global_polar : bool
+        wheter to use glovbal version of polar potential
+
+    Returns
+    -------
+    Tuple[List[np.ndarray], List[int]]
+        arrays with results and their shapes
+    """
+    data.add(
+        "polarizability" if not atomic else "atomic_polarizability",
+        9,
+        atomic=atomic,
+        must=True,
+        high_prec=False,
+        type_sel=dp.get_sel_type(),
+    )
+    
+    test_data = data.get_test()
+    polar, numb_test, atype = run_test(dp, test_data, numb_test)
+
+    sel_type = dp.get_sel_type()
+    sel_natoms = 0
+    for ii in sel_type:
+        sel_natoms += sum(atype == ii)
+
+    # YWolfeee: do summation in global polar mode
+    if not atomic:
+        polar = np.sum(polar.reshape((polar.shape[0],-1,9)),axis=1)    
+        rmse_f = rmse(polar - test_data["polarizability"][:numb_test])
+        rmse_fs = rmse_f / np.sqrt(sel_natoms)
+        rmse_fa = rmse_f / sel_natoms
+    else:
+        rmse_f = rmse(polar - test_data["atomic_polarizability"][:numb_test])
+    
+    log.info(f"# number of test data : {numb_test:d} ")
+    log.info(f"Polarizability  RMSE       : {rmse_f:e}")
+    if not atomic:
+        log.info(f"Polarizability  RMSE/sqrtN : {rmse_fs:e}")
+        log.info(f"Polarizability  RMSE/N     : {rmse_fa:e}")
+    log.info(f"The unit of error is the same as the unit of provided label.")
+
+    if detail_file is not None:
+        detail_path = Path(detail_file)
+
+        pe = np.concatenate(
+            (
+                np.reshape(test_data["polarizability"][:numb_test], [-1, 9]),
+                np.reshape(polar, [-1, 9]),
+            ),
+            axis=1,
+        )
+        np.savetxt(
+            detail_path.with_suffix(".out"),
+            pe,
+            header="data_pxx data_pxy data_pxz data_pyx data_pyy data_pyz data_pzx "
+            "data_pzy data_pzz pred_pxx pred_pxy pred_pxz pred_pyx pred_pyy pred_pyz "
+            "pred_pzx pred_pzy pred_pzz",
+        )
+    return {
+        "rmse" : (rmse_f, polar.size)
+    }
+
+
+def print_polar_sys_avg(avg):
+    """Print errors summary for polar type potential.
+
+    Parameters
+    ----------
+    avg : np.ndarray
+        array with summaries
+    """
+    log.info(f"Polarizability  RMSE : {avg['rmse']:e} eV/A")
+
+
+def test_dipole(
+    dp: "DeepDipole",
+    data: DeepmdData,
+    numb_test: int,
+    detail_file: Optional[str],
+    atomic: bool,
+) -> Tuple[List[np.ndarray], List[int]]:
+    """Test energy type model.
+
+    Parameters
+    ----------
+    dp : DeepPot
+        instance of deep potential
+    data: DeepmdData
+        data container object
+    numb_test : int
+        munber of tests to do
+    detail_file : Optional[str]
+        file where test details will be output
+    atomic : bool
+        whether atomic dipole is provided
+
+    Returns
+    -------
+    Tuple[List[np.ndarray], List[int]]
+        arrays with results and their shapes
+    """
+    data.add(
+        "dipole" if not atomic else "atomic_dipole",
+        3, 
+        atomic=atomic, 
+        must=True, 
+        high_prec=False, 
+        type_sel=dp.get_sel_type()
+    )
+    test_data = data.get_test()
+    dipole, numb_test, atype = run_test(dp, test_data, numb_test)
+
+    sel_type = dp.get_sel_type()
+    sel_natoms = 0
+    for ii in sel_type:
+        sel_natoms += sum(atype == ii)
+    
+    # do summation in atom dimension
+    if not atomic:
+        dipole = np.sum(dipole.reshape((dipole.shape[0], -1, 3)),axis=1)
+        rmse_f = rmse(dipole - test_data["dipole"][:numb_test])
+        rmse_fs = rmse_f / np.sqrt(sel_natoms)
+        rmse_fa = rmse_f / sel_natoms
+    else:
+        rmse_f = rmse(dipole - test_data["atomic_dipole"][:numb_test])
+    
+    log.info(f"# number of test data : {numb_test:d}")
+    log.info(f"Dipole  RMSE       : {rmse_f:e}")
+    if not atomic:
+        log.info(f"Dipole  RMSE/sqrtN : {rmse_fs:e}")
+        log.info(f"Dipole  RMSE/N     : {rmse_fa:e}")
+    log.info(f"The unit of error is the same as the unit of provided label.")
+
+    if detail_file is not None:
+        detail_path = Path(detail_file)
+
+        pe = np.concatenate(
+            (
+                np.reshape(test_data["dipole"][:numb_test], [-1, 3]),
+                np.reshape(dipole, [-1, 3]),
+            ),
+            axis=1,
+        )
+        np.savetxt(
+            detail_path.with_suffix(".out"),
+            pe,
+            header="data_x data_y data_z pred_x pred_y pred_z",
+        )
+    return {
+        'rmse' : (rmse_f, dipole.size)
+    }
+
+
+def print_dipole_sys_avg(avg):
+    """Print errors summary for dipole type potential.
+
+    Parameters
+    ----------
+    avg : np.ndarray
+        array with summaries
+    """
+    log.info(f"Dipole  RMSE         : {avg['rmse']:e} eV/A")

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/train.py

+"""DeePMD training entrypoint script.
+
+Can handle local or distributed training.
+"""
+
+import json
+import logging
+import time
+import os
+from typing import Dict, List, Optional, Any
+
+from deepmd.common import data_requirement, expand_sys_str, j_loader, j_must_have
+from deepmd.env import tf, reset_default_tf_session_config, GLOBAL_ENER_FLOAT_PRECISION
+from deepmd.infer.data_modifier import DipoleChargeModifier
+from deepmd.train.run_options import BUILD, CITATION, WELCOME, RunOptions
+from deepmd.train.trainer import DPTrainer
+from deepmd.utils import random as dp_random
+from deepmd.utils.argcheck import normalize
+from deepmd.utils.compat import update_deepmd_input
+from deepmd.utils.data_system import DeepmdDataSystem
+from deepmd.utils.sess import run_sess
+from deepmd.utils.neighbor_stat import NeighborStat
+from deepmd.utils.path import DPPath
+
+__all__ = ["train"]
+
+log = logging.getLogger(__name__)
+
+
+def train(
+    *,
+    INPUT: str,
+    init_model: Optional[str],
+    restart: Optional[str],
+    output: str,
+    init_frz_model: str,
+    mpi_log: str,
+    log_level: int,
+    log_path: Optional[str],
+    is_compress: bool = False,
+    skip_neighbor_stat: bool = False,
+    **kwargs,
+):
+    """Run DeePMD model training.
+
+    Parameters
+    ----------
+    INPUT : str
+        json/yaml control file
+    init_model : Optional[str]
+        path to checkpoint folder or None
+    restart : Optional[str]
+        path to checkpoint folder or None
+    output : str
+        path for dump file with arguments
+    init_frz_model : str
+        path to frozen model or None
+    mpi_log : str
+        mpi logging mode
+    log_level : int
+        logging level defined by int 0-3
+    log_path : Optional[str]
+        logging file path or None if logs are to be output only to stdout
+    is_compress: bool
+        indicates whether in the model compress mode
+    skip_neighbor_stat : bool, default=False
+        skip checking neighbor statistics
+
+    Raises
+    ------
+    RuntimeError
+        if distributed training job nem is wrong
+    """
+    run_opt = RunOptions(
+        init_model=init_model,
+        restart=restart,
+        init_frz_model=init_frz_model,
+        log_path=log_path,
+        log_level=log_level,
+        mpi_log=mpi_log
+    )
+    if run_opt.is_distrib and len(run_opt.gpus or []) > 1:
+        # avoid conflict of visible gpus among multipe tf sessions in one process
+        reset_default_tf_session_config(cpu_only=True)
+
+    # load json database
+    jdata = j_loader(INPUT)
+
+    jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json")
+
+    jdata = normalize(jdata)
+
+    if not is_compress and not skip_neighbor_stat:
+        jdata = update_sel(jdata)
+
+    with open(output, "w") as fp:
+        json.dump(jdata, fp, indent=4)
+
+    # save the training script into the graph
+    # remove white spaces as it is not compressed
+    tf.constant(json.dumps(jdata, separators=(',', ':')), name='train_attr/training_script', dtype=tf.string)
+
+    for message in WELCOME + CITATION + BUILD:
+        log.info(message)
+
+    run_opt.print_resource_summary()
+    _do_work(jdata, run_opt, is_compress)
+
+
+def _do_work(jdata: Dict[str, Any], run_opt: RunOptions, is_compress: bool = False):
+    """Run serial model training.
+
+    Parameters
+    ----------
+    jdata : Dict[str, Any]
+        arguments read form json/yaml control file
+    run_opt : RunOptions
+        object with run configuration
+    is_compress : Bool
+        indicates whether in model compress mode
+
+    Raises
+    ------
+    RuntimeError
+        If unsupported modifier type is selected for model
+    """
+    # make necessary checks
+    assert "training" in jdata
+
+    # init the model
+    model = DPTrainer(jdata, run_opt=run_opt, is_compress = is_compress)
+    rcut = model.model.get_rcut()
+    type_map = model.model.get_type_map()
+    if len(type_map) == 0:
+        ipt_type_map = None
+    else:
+        ipt_type_map = type_map
+
+    # init random seed of data systems
+    seed = jdata["training"].get("seed", None)
+    if seed is not None:
+        # avoid the same batch sequence among workers
+        seed += run_opt.my_rank
+        seed = seed % (2 ** 32)
+    dp_random.seed(seed)
+
+    # setup data modifier
+    modifier = get_modifier(jdata["model"].get("modifier", None))
+
+    # decouple the training data from the model compress process
+    train_data = None
+    valid_data = None
+    if not is_compress:
+        # init data
+        train_data = get_data(jdata["training"]["training_data"], rcut, ipt_type_map, modifier)
+        train_data.print_summary("training")
+        if jdata["training"].get("validation_data", None) is not None:
+            valid_data = get_data(jdata["training"]["validation_data"], rcut, train_data.type_map, modifier)
+            valid_data.print_summary("validation")
+
+    # get training info
+    stop_batch = j_must_have(jdata["training"], "numb_steps")
+    model.build(train_data, stop_batch)
+
+    if not is_compress:
+        # train the model with the provided systems in a cyclic way
+        start_time = time.time()
+        model.train(train_data, valid_data)
+        end_time = time.time()
+        log.info("finished training")
+        log.info(f"wall time: {(end_time - start_time):.3f} s")
+    else:
+        model.save_compressed()
+        log.info("finished compressing")
+
+
+def get_data(jdata: Dict[str, Any], rcut, type_map, modifier):
+    systems = j_must_have(jdata, "systems")
+    if isinstance(systems, str):
+        systems = expand_sys_str(systems)
+    help_msg = 'Please check your setting for data systems'
+    # check length of systems
+    if len(systems) == 0:
+        msg = 'cannot find valid a data system'
+        log.fatal(msg)
+        raise IOError(msg, help_msg)
+    # rougly check all items in systems are valid
+    for ii in systems:
+        ii = DPPath(ii)
+        if (not ii.is_dir()):
+            msg = f'dir {ii} is not a valid dir'
+            log.fatal(msg)
+            raise IOError(msg, help_msg)
+        if (not (ii / 'type.raw').is_file()):
+            msg = f'dir {ii} is not a valid data system dir'
+            log.fatal(msg)
+            raise IOError(msg, help_msg)
+
+    batch_size = j_must_have(jdata, "batch_size")
+    sys_probs = jdata.get("sys_probs", None)
+    auto_prob = jdata.get("auto_prob", "prob_sys_size")
+
+    data = DeepmdDataSystem(
+        systems=systems,
+        batch_size=batch_size,
+        test_size=1,        # to satisfy the old api
+        shuffle_test=True,  # to satisfy the old api
+        rcut=rcut,
+        type_map=type_map,
+        modifier=modifier,
+        trn_all_set=True,    # sample from all sets
+        sys_probs=sys_probs,
+        auto_prob_style=auto_prob
+    )
+    data.add_dict(data_requirement)
+
+    return data
+
+
+def get_modifier(modi_data=None):
+    modifier: Optional[DipoleChargeModifier]
+    if modi_data is not None:
+        if modi_data["type"] == "dipole_charge":
+            modifier = DipoleChargeModifier(
+                modi_data["model_name"],
+                modi_data["model_charge_map"],
+                modi_data["sys_charge_map"],
+                modi_data["ewald_h"],
+                modi_data["ewald_beta"],
+            )
+        else:
+            raise RuntimeError("unknown modifier type " + str(modi_data["type"]))
+    else:
+        modifier = None
+    return modifier
+
+
+def get_rcut(jdata):
+    descrpt_data = jdata['model']['descriptor']
+    rcut_list = []
+    if descrpt_data['type'] == 'hybrid':
+        for ii in descrpt_data['list']:
+            rcut_list.append(ii['rcut'])
+    else:
+        rcut_list.append(descrpt_data['rcut'])
+    return max(rcut_list)
+
+
+def get_type_map(jdata):
+    return jdata['model'].get('type_map', None)
+
+
+def get_nbor_stat(jdata, rcut, one_type: bool = False):
+    max_rcut = get_rcut(jdata)
+    type_map = get_type_map(jdata)
+
+    if type_map and len(type_map) == 0:
+        type_map = None
+    train_data = get_data(jdata["training"]["training_data"], max_rcut, type_map, None)
+    train_data.get_batch()
+    data_ntypes = train_data.get_ntypes()
+    if type_map is not None:
+        map_ntypes = len(type_map)
+    else:
+        map_ntypes = data_ntypes
+    ntypes = max([map_ntypes, data_ntypes])
+
+    neistat = NeighborStat(ntypes, rcut, one_type=one_type)
+
+    min_nbor_dist, max_nbor_size = neistat.get_stat(train_data)
+
+    # moved from traier.py as duplicated
+    # TODO: this is a simple fix but we should have a clear
+    #       architecture to call neighbor stat
+    tf.constant(min_nbor_dist,
+        name = 'train_attr/min_nbor_dist',
+        dtype = GLOBAL_ENER_FLOAT_PRECISION)
+    tf.constant(max_nbor_size,
+        name = 'train_attr/max_nbor_size',
+        dtype = tf.int32)
+    return min_nbor_dist, max_nbor_size
+
+def get_sel(jdata, rcut, one_type: bool = False):
+    _, max_nbor_size = get_nbor_stat(jdata, rcut, one_type=one_type)
+    return max_nbor_size
+
+def get_min_nbor_dist(jdata, rcut):
+    min_nbor_dist, _ = get_nbor_stat(jdata, rcut)
+    return min_nbor_dist
+
+def parse_auto_sel(sel):
+    if type(sel) is not str:
+        return False
+    words = sel.split(':')
+    if words[0] == 'auto':
+        return True
+    else:
+        return False
+
+    
+def parse_auto_sel_ratio(sel):
+    if not parse_auto_sel(sel):
+        raise RuntimeError(f'invalid auto sel format {sel}')
+    else:
+        words = sel.split(':')
+        if len(words) == 1:
+            ratio = 1.1
+        elif len(words) == 2:
+            ratio = float(words[1])
+        else:
+            raise RuntimeError(f'invalid auto sel format {sel}')
+        return ratio
+
+
+def wrap_up_4(xx):
+    return 4 * ((int(xx) + 3) // 4)
+
+
+def update_one_sel(jdata, descriptor):
+    if descriptor['type'] == 'loc_frame':
+        return descriptor
+    rcut = descriptor['rcut']
+    tmp_sel = get_sel(jdata, rcut, one_type=descriptor['type'] in ('se_atten',))
+    sel = descriptor['sel']
+    if isinstance(sel, int):
+        # convert to list and finnally convert back to int
+        sel = [sel]
+    if parse_auto_sel(descriptor['sel']) :
+        ratio = parse_auto_sel_ratio(descriptor['sel'])
+        descriptor['sel'] = sel = [int(wrap_up_4(ii * ratio)) for ii in tmp_sel]
+    else:
+        # sel is set by user
+        for ii, (tt, dd) in enumerate(zip(tmp_sel, sel)):
+            if dd and tt > dd:
+                # we may skip warning for sel=0, where the user is likely
+                # to exclude such type in the descriptor
+                log.warning(
+                    "sel of type %d is not enough! The expected value is "
+                    "not less than %d, but you set it to %d. The accuracy"
+                    " of your model may get worse." %(ii, tt, dd)
+                )
+    if descriptor['type'] in ('se_atten',):
+        descriptor['sel'] = sel = sum(sel)
+    return descriptor
+
+
+def update_sel(jdata):    
+    log.info("Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)")
+    descrpt_data = jdata['model']['descriptor']
+    if descrpt_data['type'] == 'hybrid':
+        for ii in range(len(descrpt_data['list'])):
+            descrpt_data['list'][ii] = update_one_sel(jdata, descrpt_data['list'][ii])
+    else:
+        descrpt_data = update_one_sel(jdata, descrpt_data)
+    jdata['model']['descriptor'] = descrpt_data
+    return jdata

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/entrypoints/transfer.py

+"""Module used for transfering parameters between models."""
+
+from typing import Dict, Optional, Sequence, Tuple
+from deepmd.env import tf, TRANSFER_PATTERN
+import re
+import numpy as np
+import logging
+
+__all__ = ["transfer"]
+
+log = logging.getLogger(__name__)
+
+
+@np.vectorize
+def convert_number(number: int) -> float:
+    binary = bin(number).replace("0b", "").zfill(16)
+    sign = int(binary[0]) * -2 + 1
+    exp = int(binary[1:6], 2)
+    frac = (int(binary[6:], 2) + 2 ** 10) * (2 ** -10)
+    return sign * (2 ** (exp - 15)) * frac
+
+
+def convert_matrix(
+    matrix: np.ndarray, shape: Sequence[int], dtype: Optional[type] = None
+) -> np.ndarray:
+    """Convert matrix of integers to self defined binary format.
+
+    Parameters
+    ----------
+    matrix : np.ndarray
+        array of ints
+    shape : Sequence[int]
+        shape to cast resulting array to
+    dtype : Optional[type]
+        type that finall array will be cast to, If None no casting will take place
+
+    Returns
+    -------
+    np.ndarray
+        array cast to required format
+    """
+    conv = convert_number(matrix.flatten()).reshape(shape)
+    if dtype:
+        conv = conv.astype(dtype)
+
+    return conv
+
+
+def transfer(*, old_model: str, raw_model: str, output: str, **kwargs):
+    """Transfer operation from old fron graph to new prepared raw graph.
+
+    Parameters
+    ----------
+    old_model : str
+        frozen old graph model
+    raw_model : str
+        new model that will accept ops from old model
+    output : str
+        new model with transfered parameters will be saved to this location
+    """
+    raw_graph = load_graph(raw_model)
+    old_graph = load_graph(old_model)
+    log.info(f"{len(raw_graph.as_graph_def().node)} ops in the raw graph")
+    log.info(f"{len(old_graph.as_graph_def().node)} ops in the old graph")
+
+    new_graph_def = transform_graph(raw_graph, old_graph)
+    with tf.gfile.GFile(output, mode="wb") as f:
+        f.write(new_graph_def.SerializeToString())
+    log.info("the output model is saved in " + output)
+
+
+def load_graph(graph_name: str) -> tf.Graph:
+    """Load graph from passed in path.
+
+    Parameters
+    ----------
+    graph_name : str
+        path to frozen graph on disk
+
+    Returns
+    -------
+    tf.Graph
+        tf graph object
+    """
+    graph_def = tf.GraphDef()
+    with open(graph_name, "rb") as f:
+        graph_def.ParseFromString(f.read())
+    with tf.Graph().as_default() as graph:
+        tf.import_graph_def(graph_def, name="")
+        return graph
+
+
+def transform_graph(raw_graph: tf.Graph, old_graph: tf.Graph) -> tf.Graph:
+    """Trasform old graph into new.
+
+    Parameters
+    ----------
+    raw_graph : tf.Graph
+        graph receiving parameters from the old one
+    old_graph : tf.Graph
+        graph providing parameters
+
+    Returns
+    -------
+    tf.Graph
+        new graph with parameters transfered form the old one
+    """
+    old_graph_def = old_graph.as_graph_def()
+    raw_graph_def = raw_graph.as_graph_def()
+    raw_graph_node = load_transform_node(raw_graph_def)
+    old_graph_node = load_transform_node(old_graph_def)
+
+    for node in raw_graph_def.node:
+        if node.name not in raw_graph_node.keys():
+            continue
+
+        old_node = old_graph_node[node.name]
+        raw_node = raw_graph_node[node.name]
+        cp_attr = CopyNodeAttr(node)
+
+        check_dim(raw_graph_node, old_graph_node, node.name)
+        tensor_shape = [dim.size for dim in raw_node.tensor_shape.dim]
+        old_graph_dtype = tf.as_dtype(old_node.dtype).as_numpy_dtype
+        raw_graph_dtype = tf.as_dtype(raw_node.dtype).as_numpy_dtype
+        log.info(
+            f"{node.name} is passed from old graph({old_graph_dtype}) "
+            f"to raw graph({raw_graph_dtype})"
+        )
+
+        if raw_graph_dtype == np.float16:
+            if old_graph_dtype == np.float64 or old_graph_dtype == np.float32:
+                if (len(tensor_shape) != 1) or (tensor_shape[0] != 1):
+                    tensor = np.frombuffer(old_node.tensor_content, dtype = old_graph_dtype)
+                    tensor = tensor.astype(raw_graph_dtype)
+                    cp_attr.from_str(tensor)
+                else:
+                    tensor = load_tensor(old_node, old_graph_dtype, raw_graph_dtype)
+                    cp_attr.from_array(tensor, tf.float16, [1])
+
+            elif old_graph_dtype[1] == "float16":
+                tensor = convertMatrix(np.array(old_node.half_val), tensor_shape)
+                cp_attr.from_array(tensor, raw_graph_dtype)
+
+        elif raw_graph_dtype == np.float64 or raw_graph_dtype == np.float32:
+            if old_graph_dtype == np.float64 or old_graph_dtype == np.float32:
+                if (len(tensor_shape) != 1) or (tensor_shape[0] != 1):
+                    tensor = np.frombuffer(old_node.tensor_content, dtype = old_graph_dtype)
+                    tensor = tensor.astype(raw_graph_dtype)
+                    cp_attr.from_str(tensor)
+                else:
+                    tensor = load_tensor(old_node, old_graph_dtype, raw_graph_dtype)
+                    cp_attr.from_array(tensor, raw_graph_dtype, shape=[1])
+
+            elif old_graph_dtype == np.float16:
+                if (len(tensor_shape) != 1) or (tensor_shape[0] != 1):
+                    tensor = convertMatrix(np.array(old_node.half_val), tensor_shape).astype(raw_graph_dtype)
+                    cp_attr.from_str(tensor)
+                else:
+                    tensor = convertMatrix(np.array(old_node.half_val), tensor_shape).astype(raw_graph_dtype)
+                    cp_attr.from_array(tensor, raw_graph_dtype)
+
+    return raw_graph_def
+
+
+class CopyNodeAttr:
+    def __init__(self, node) -> None:
+        self.node = node
+
+    def from_array(
+        self, tensor: np.ndarray, dtype: type, shape: Optional[Sequence[int]] = None
+    ):
+        if shape is None:
+            shape = tensor.shape
+        self.node.attr["value"].CopyFrom(
+            tf.AttrValue(tensor=tf.make_tensor_proto(tensor, dtype, shape))
+        )
+
+    def from_str(self, tensor: np.ndarray):
+        self.node.attr["value"].tensor.tensor_content = tensor.tostring()
+
+
+def load_tensor(node: tf.Tensor, dtype_old: type, dtype_new: type) -> np.ndarray:
+    if dtype_old == np.float64:
+        tensor = np.array(node.double_val).astype(dtype_new)
+    elif dtype_old == np.float32:
+        tensor = np.array(node.float_val).astype(dtype_new)
+
+    return tensor
+
+
+def check_dim(raw_graph_node: tf.Tensor, old_graph_node: tf.Tensor, node_name: str):
+    """Check if dimensions of tensor in old and new graph is equal.
+
+    Parameters
+    ----------
+    raw_graph_node : tf.Tensor
+        node of the receiving graph
+    old_graph_node : tf.Tensor
+        node of the graph from which will node be extracted
+    node_name : str
+        name of the node
+
+    Raises
+    ------
+    RuntimeError
+        if node dimension do not match
+    """
+    raw_graph_dim = raw_graph_node[node_name].tensor_shape
+    old_graph_dim = old_graph_node[node_name].tensor_shape
+    if raw_graph_dim != old_graph_dim:
+        raise RuntimeError(
+            f"old graph {old_graph_dim} and raw graph {raw_graph_dim} "
+            f"has different {node_name} dim"
+        )
+
+
+def load_transform_node(graph: tf.Graph) -> Dict[str, tf.Tensor]:
+    """Load nodes and their names from graph to dict.
+
+    Parameters
+    ----------
+    graph : tf.Graph
+        tensforflow graph
+
+    Returns
+    -------
+    Dict[str, tf.Tensor]
+        mapping on graph node names and corresponding tensors
+    """
+    transform_node_pattern = re.compile(TRANSFER_PATTERN)
+
+    transform_node = {}
+    for node in graph.node:
+        if transform_node_pattern.fullmatch(node.name) is not None:
+            transform_node[node.name] = node.attr["value"].tensor
+    return transform_node

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/env.py

+"""Module that sets tensorflow working environment and exports inportant constants."""
+
+import logging
+import os
+import re
+import platform
+from configparser import ConfigParser
+from importlib import reload
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+from packaging.version import Version
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from types import ModuleType
+
+# import tensorflow v1 compatability
+try:
+    import tensorflow.compat.v1 as tf
+
+    tf.disable_v2_behavior()
+except ImportError:
+    import tensorflow as tf
+try:
+    import tensorflow.compat.v2 as tfv2
+except ImportError:
+    tfv2 = None
+
+__all__ = [
+    "GLOBAL_CONFIG",
+    "GLOBAL_TF_FLOAT_PRECISION",
+    "GLOBAL_NP_FLOAT_PRECISION",
+    "GLOBAL_ENER_FLOAT_PRECISION",
+    "global_float_prec",
+    "global_cvt_2_tf_float",
+    "global_cvt_2_ener_float",
+    "MODEL_VERSION",
+    "SHARED_LIB_MODULE",
+    "default_tf_session_config",
+    "reset_default_tf_session_config",
+    "op_module",
+    "op_grads_module",
+    "TRANSFER_PATTERN",
+    "FITTING_NET_PATTERN",
+    "EMBEDDING_NET_PATTERN",
+    "TYPE_EMBEDDING_PATTERN",
+    "ATTENTION_LAYER_PATTERN",
+    "TF_VERSION"
+]
+
+SHARED_LIB_MODULE = "op"
+
+# Python library version
+try:
+    tf_py_version = tf.version.VERSION
+except AttributeError:
+    tf_py_version = tf.__version__
+
+EMBEDDING_NET_PATTERN = str(
+    r"filter_type_\d+/matrix_\d+_\d+|"
+    r"filter_type_\d+/bias_\d+_\d+|"
+    r"filter_type_\d+/idt_\d+_\d+|"
+    r"filter_type_all/matrix_\d+|"
+    r"filter_type_all/matrix_\d+_\d+|"
+    r"filter_type_all/matrix_\d+_\d+_\d+|"
+    r"filter_type_all/bias_\d+|"
+    r"filter_type_all/bias_\d+_\d+|"
+    r"filter_type_all/bias_\d+_\d+_\d+|"
+    r"filter_type_all/idt_\d+|"
+    r"filter_type_all/idt_\d+_\d+|"
+)
+
+FITTING_NET_PATTERN = str(
+    r"layer_\d+/matrix|"
+    r"layer_\d+_type_\d+/matrix|"
+    r"layer_\d+/bias|"
+    r"layer_\d+_type_\d+/bias|"
+    r"layer_\d+/idt|"
+    r"layer_\d+_type_\d+/idt|"
+    r"final_layer/matrix|"
+    r"final_layer_type_\d+/matrix|"
+    r"final_layer/bias|"
+    r"final_layer_type_\d+/bias|"
+)
+
+TYPE_EMBEDDING_PATTERN = str(
+    r"type_embed_net+/matrix_\d+|"
+    r"type_embed_net+/bias_\d+|"
+    r"type_embed_net+/idt_\d+|"
+)
+
+ATTENTION_LAYER_PATTERN = str(
+    r"attention_layer_\d+/c_query/matrix|"
+    r"attention_layer_\d+/c_query/bias|"
+    r"attention_layer_\d+/c_key/matrix|"
+    r"attention_layer_\d+/c_key/bias|"
+    r"attention_layer_\d+/c_value/matrix|"
+    r"attention_layer_\d+/c_value/bias|"
+    r"attention_layer_\d+/c_out/matrix|"
+    r"attention_layer_\d+/c_out/bias|"
+    r"attention_layer_\d+/layer_normalization/beta|"
+    r"attention_layer_\d+/layer_normalization/gamma|"
+    r"attention_layer_\d+/layer_normalization_\d+/beta|"
+    r"attention_layer_\d+/layer_normalization_\d+/gamma|"
+)
+
+TRANSFER_PATTERN = \
+    EMBEDDING_NET_PATTERN + \
+    FITTING_NET_PATTERN + \
+    TYPE_EMBEDDING_PATTERN + \
+    str(
+        r"descrpt_attr/t_avg|"
+        r"descrpt_attr/t_std|"
+        r"fitting_attr/t_fparam_avg|"
+        r"fitting_attr/t_fparam_istd|"
+        r"fitting_attr/t_aparam_avg|"
+        r"fitting_attr/t_aparam_istd|"
+        r"model_attr/t_tab_info|"
+        r"model_attr/t_tab_data|"
+)
+
+def set_env_if_empty(key: str, value: str, verbose: bool = True):
+    """Set environment variable only if it is empty.
+
+    Parameters
+    ----------
+    key : str
+        env variable name
+    value : str
+        env variable value
+    verbose : bool, optional
+        if True action will be logged, by default True
+    """
+    if os.environ.get(key) is None:
+        os.environ[key] = value
+        if verbose:
+            logging.warn(
+                f"Environment variable {key} is empty. Use the default value {value}"
+            )
+
+
+def set_mkl():
+    """Tuning MKL for the best performance.
+
+    References
+    ----------
+    TF overview
+    https://www.tensorflow.org/guide/performance/overview
+
+    Fixing an issue in numpy built by MKL
+    https://github.com/ContinuumIO/anaconda-issues/issues/11367
+    https://github.com/numpy/numpy/issues/12374
+
+    check whether the numpy is built by mkl, see
+    https://github.com/numpy/numpy/issues/14751
+    """
+    if "mkl_rt" in np.__config__.get_info("blas_mkl_info").get("libraries", []):
+        set_env_if_empty("KMP_BLOCKTIME", "0")
+        set_env_if_empty(
+            "KMP_AFFINITY", "granularity=fine,verbose,compact,1,0")
+        reload(np)
+
+
+def set_tf_default_nthreads():
+    """Set TF internal number of threads to default=automatic selection.
+
+    Notes
+    -----
+    `TF_INTRA_OP_PARALLELISM_THREADS` and `TF_INTER_OP_PARALLELISM_THREADS`
+    control TF configuration of multithreading.
+    """
+    if "OMP_NUM_THREADS" not in os.environ or \
+       "TF_INTRA_OP_PARALLELISM_THREADS" not in os.environ or \
+       "TF_INTER_OP_PARALLELISM_THREADS" not in os.environ:
+        logging.warning(
+            "To get the best performance, it is recommended to adjust "
+            "the number of threads by setting the environment variables "
+            "OMP_NUM_THREADS, TF_INTRA_OP_PARALLELISM_THREADS, and "
+            "TF_INTER_OP_PARALLELISM_THREADS.")
+    set_env_if_empty("TF_INTRA_OP_PARALLELISM_THREADS", "0", verbose=False)
+    set_env_if_empty("TF_INTER_OP_PARALLELISM_THREADS", "0", verbose=False)
+
+
+def get_tf_default_nthreads() -> Tuple[int, int]:
+    """Get TF paralellism settings.
+
+    Returns
+    -------
+    Tuple[int, int]
+        number of `TF_INTRA_OP_PARALLELISM_THREADS` and
+        `TF_INTER_OP_PARALLELISM_THREADS`
+    """
+    return int(os.environ.get("TF_INTRA_OP_PARALLELISM_THREADS", "0")), int(
+        os.environ.get("TF_INTER_OP_PARALLELISM_THREADS", "0")
+    )
+
+
+def get_tf_session_config() -> Any:
+    """Configure tensorflow session.
+
+    Returns
+    -------
+    Any
+        session configure object
+    """
+    set_tf_default_nthreads()
+    intra, inter = get_tf_default_nthreads()
+    config = tf.ConfigProto(
+        gpu_options=tf.GPUOptions(allow_growth=True),
+        intra_op_parallelism_threads=intra, inter_op_parallelism_threads=inter
+    )
+    if Version(tf_py_version) >= Version('1.15') and int(os.environ.get("DP_AUTO_PARALLELIZATION", 0)):
+        config.graph_options.rewrite_options.custom_optimizers.add().name = "dpparallel"
+    return config
+
+
+default_tf_session_config = get_tf_session_config()
+
+
+def reset_default_tf_session_config(cpu_only: bool):
+    """Limit tensorflow session to CPU or not.
+
+    Parameters
+    ----------
+    cpu_only : bool
+        If enabled, no GPU device is visible to the TensorFlow Session.
+    """
+    global default_tf_session_config
+    if cpu_only:
+        default_tf_session_config.device_count['GPU'] = 0
+    else:
+        if 'GPU' in default_tf_session_config.device_count:
+            del default_tf_session_config.device_count['GPU']
+
+
+def get_module(module_name: str) -> "ModuleType":
+    """Load force module.
+
+    Returns
+    -------
+    ModuleType
+        loaded force module
+
+    Raises
+    ------
+    FileNotFoundError
+        if module is not found in directory
+    """
+    if platform.system() == "Windows":
+        ext = ".dll"
+        prefix = ""
+    #elif platform.system() == "Darwin":
+    #    ext = ".dylib"
+    else:
+        ext = ".so"
+        prefix = "lib"
+
+    module_file = (
+        (Path(__file__).parent / SHARED_LIB_MODULE / (prefix + module_name))
+        .with_suffix(ext)
+        .resolve()
+    )
+
+    if not module_file.is_file():
+        raise FileNotFoundError(f"module {module_name} does not exist")
+    else:
+        try:
+            module = tf.load_op_library(str(module_file))
+        except tf.errors.NotFoundError as e:
+            # check CXX11_ABI_FLAG is compatiblity
+            # see https://gcc.gnu.org/onlinedocs/libstdc++/manual/using_dual_abi.html
+            # ABI should be the same
+            if 'CXX11_ABI_FLAG' in tf.__dict__:
+                tf_cxx11_abi_flag = tf.CXX11_ABI_FLAG
+            else:
+                tf_cxx11_abi_flag = tf.sysconfig.CXX11_ABI_FLAG
+            if TF_CXX11_ABI_FLAG != tf_cxx11_abi_flag:
+                raise RuntimeError(
+                    "This deepmd-kit package was compiled with "
+                    "CXX11_ABI_FLAG=%d, but TensorFlow runtime was compiled "
+                    "with CXX11_ABI_FLAG=%d. These two library ABIs are "
+                    "incompatible and thus an error is raised when loading %s. "
+                    "You need to rebuild deepmd-kit against this TensorFlow "
+                    "runtime." % (
+                        TF_CXX11_ABI_FLAG,
+                        tf_cxx11_abi_flag,
+                        module_name,
+                    )) from e
+
+            # different versions may cause incompatibility
+            # see #406, #447, #557, #774, and #796 for example
+            # throw a message if versions are different
+            if TF_VERSION != tf_py_version:
+                raise RuntimeError(
+                    "The version of TensorFlow used to compile this "
+                    "deepmd-kit package is %s, but the version of TensorFlow "
+                    "runtime you are using is %s. These two versions are "
+                    "incompatible and thus an error is raised when loading %s. "
+                    "You need to install TensorFlow %s, or rebuild deepmd-kit "
+                    "against TensorFlow %s.\nIf you are using a wheel from "
+                    "pypi, you may consider to install deepmd-kit execuating "
+                    "`pip install deepmd-kit --no-binary deepmd-kit` "
+                    "instead." % (
+                        TF_VERSION,
+                        tf_py_version,
+                        module_name,
+                        TF_VERSION,
+                        tf_py_version,
+                    )) from e
+            error_message = (
+                "This deepmd-kit package is inconsitent with TensorFlow "
+                "Runtime, thus an error is raised when loading %s. "
+                "You need to rebuild deepmd-kit against this TensorFlow "
+                "runtime." % (
+                    module_name,
+                )
+            )
+            if TF_CXX11_ABI_FLAG == 1:
+                # #1791
+                error_message += (
+                    "\nWARNING: devtoolset on RHEL6 and RHEL7 does not support _GLIBCXX_USE_CXX11_ABI=1. "
+                    "See https://bugzilla.redhat.com/show_bug.cgi?id=1546704"
+                )
+            raise RuntimeError(error_message) from e
+        return module
+
+
+def _get_package_constants(
+    config_file: Path = Path(__file__).parent / "pkg_config/run_config.ini",
+) -> Dict[str, str]:
+    """Read package constants set at compile time by CMake to dictionary.
+
+    Parameters
+    ----------
+    config_file : str, optional
+        path to CONFIG file, by default "pkg_config/run_config.ini"
+
+    Returns
+    -------
+    Dict[str, str]
+        dictionary with package constants
+    """
+    config = ConfigParser()
+    config.read(config_file)
+    return dict(config.items("CONFIG"))
+
+
+GLOBAL_CONFIG = _get_package_constants()
+MODEL_VERSION = GLOBAL_CONFIG["model_version"]
+TF_VERSION = GLOBAL_CONFIG["tf_version"]
+TF_CXX11_ABI_FLAG = int(GLOBAL_CONFIG["tf_cxx11_abi_flag"])
+
+op_module = get_module("op_abi")
+op_grads_module = get_module("op_grads")
+
+# FLOAT_PREC
+dp_float_prec = os.environ.get("DP_INTERFACE_PREC", "high").lower()
+if dp_float_prec in ("high", ""):
+    # default is high
+    GLOBAL_TF_FLOAT_PRECISION = tf.float64
+    GLOBAL_NP_FLOAT_PRECISION = np.float64
+    GLOBAL_ENER_FLOAT_PRECISION = np.float64
+    global_float_prec = "double"
+elif dp_float_prec == "low":
+    GLOBAL_TF_FLOAT_PRECISION = tf.float32
+    GLOBAL_NP_FLOAT_PRECISION = np.float32
+    GLOBAL_ENER_FLOAT_PRECISION = np.float64
+    global_float_prec = "float"
+else:
+    raise RuntimeError(
+        "Unsupported float precision option: %s. Supported: high,"
+        "low. Please set precision with environmental variable "
+        "DP_INTERFACE_PREC." % dp_float_prec
+    )
+
+
+def global_cvt_2_tf_float(xx: tf.Tensor) -> tf.Tensor:
+    """Cast tensor to globally set TF precision.
+
+    Parameters
+    ----------
+    xx : tf.Tensor
+        input tensor
+
+    Returns
+    -------
+    tf.Tensor
+        output tensor cast to `GLOBAL_TF_FLOAT_PRECISION`
+    """
+    return tf.cast(xx, GLOBAL_TF_FLOAT_PRECISION)
+
+
+def global_cvt_2_ener_float(xx: tf.Tensor) -> tf.Tensor:
+    """Cast tensor to globally set energy precision.
+
+    Parameters
+    ----------
+    xx : tf.Tensor
+        input tensor
+
+    Returns
+    -------
+    tf.Tensor
+        output tensor cast to `GLOBAL_ENER_FLOAT_PRECISION`
+    """
+    return tf.cast(xx, GLOBAL_ENER_FLOAT_PRECISION)

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/fit/__init__.py

+from .ener      import EnerFitting
+from .wfc       import WFCFitting
+from .dipole    import DipoleFittingSeA
+from .polar     import PolarFittingSeA
+from .polar     import GlobalPolarFittingSeA
+from .polar     import PolarFittingLocFrame

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/fit/dipole.py

+import warnings
+import numpy as np
+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.common import add_data_requirement, get_activation_func, get_precision, cast_precision
+from deepmd.utils.network import one_layer, one_layer_rand_seed_shift
+from deepmd.utils.graph import get_fitting_net_variables_from_graph_def
+from deepmd.descriptor import DescrptSeA
+from deepmd.fit.fitting import Fitting
+
+from deepmd.env import global_cvt_2_tf_float
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+
+class DipoleFittingSeA (Fitting) :
+    """
+    Fit the atomic dipole with descriptor se_a
+    
+    Parameters
+    ----------
+    descrpt : tf.Tensor
+            The descrptor
+    neuron : List[int]
+            Number of neurons in each hidden layer of the fitting net
+    resnet_dt : bool
+            Time-step `dt` in the resnet construction:
+            y = x + dt * \phi (Wx + b)
+    sel_type : List[int]
+            The atom types selected to have an atomic dipole prediction. If is None, all atoms are selected.
+    seed : int
+            Random seed for initializing the network parameters.
+    activation_function : str
+            The activation function in the embedding net. Supported options are |ACTIVATION_FN|
+    precision : str
+            The precision of the embedding net parameters. Supported options are |PRECISION|
+    uniform_seed
+            Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
+    """
+    def __init__ (self, 
+                  descrpt : tf.Tensor,
+                  neuron : List[int] = [120,120,120], 
+                  resnet_dt : bool = True,
+                  sel_type : List[int] = None,
+                  seed : int = None,
+                  activation_function : str = 'tanh',
+                  precision : str = 'default',
+                  uniform_seed: bool = False
+    ) -> None:
+        """
+        Constructor
+        """
+        if not isinstance(descrpt, DescrptSeA) :
+            raise RuntimeError('DipoleFittingSeA only supports DescrptSeA')
+        self.ntypes = descrpt.get_ntypes()
+        self.dim_descrpt = descrpt.get_dim_out()
+        # args = ClassArg()\
+        #        .add('neuron',           list,   default = [120,120,120], alias = 'n_neuron')\
+        #        .add('resnet_dt',        bool,   default = True)\
+        #        .add('sel_type',         [list,int],   default = [ii for ii in range(self.ntypes)], alias = 'dipole_type')\
+        #        .add('seed',             int)\
+        #        .add("activation_function", str, default = "tanh")\
+        #        .add('precision',           str,    default = "default")
+        # class_data = args.parse(jdata)
+        self.n_neuron = neuron
+        self.resnet_dt = resnet_dt
+        self.sel_type = sel_type
+        if self.sel_type is None:
+            self.sel_type = [ii for ii in range(self.ntypes)]
+        self.seed = seed
+        self.uniform_seed = uniform_seed
+        self.seed_shift = one_layer_rand_seed_shift()
+        self.fitting_activation_fn = get_activation_func(activation_function)
+        self.fitting_precision = get_precision(precision)
+        self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1()
+        self.dim_rot_mat = self.dim_rot_mat_1 * 3
+        self.useBN = False
+        self.fitting_net_variables = None
+        self.mixed_prec = None
+
+    def get_sel_type(self) -> int:
+        """
+        Get selected type
+        """
+        return self.sel_type
+
+    def get_out_size(self) -> int:
+        """
+        Get the output size. Should be 3
+        """
+        return 3
+
+    @cast_precision
+    def build (self, 
+               input_d : tf.Tensor,
+               rot_mat : tf.Tensor,
+               natoms : tf.Tensor,
+               reuse : bool = None,
+               suffix : str = '') -> tf.Tensor:
+        """
+        Build the computational graph for fitting net
+        
+        Parameters
+        ----------
+        input_d
+                The input descriptor
+        rot_mat
+                The rotation matrix from the descriptor.
+        natoms
+                The number of atoms. This tensor has the length of Ntypes + 2
+                natoms[0]: number of local atoms
+                natoms[1]: total number of atoms held by this processor
+                natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
+        reuse
+                The weights in the networks should be reused when get the variable.
+        suffix
+                Name suffix to identify this descriptor
+
+        Returns
+        -------
+        dipole
+                The atomic dipole.
+        """
+        start_index = 0
+        inputs = tf.reshape(input_d, [-1, natoms[0], self.dim_descrpt])
+        rot_mat = tf.reshape(rot_mat, [-1, natoms[0], self.dim_rot_mat])
+
+        count = 0
+        outs_list = []
+        for type_i in range(self.ntypes):
+            # cut-out inputs
+            inputs_i = tf.slice (inputs,
+                                 [ 0, start_index, 0],
+                                 [-1, natoms[2+type_i], -1] )
+            inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
+            rot_mat_i = tf.slice (rot_mat,
+                                  [ 0, start_index, 0],
+                                  [-1, natoms[2+type_i], -1] )
+            rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3])
+            start_index += natoms[2+type_i]
+            if not type_i in self.sel_type :
+                continue
+            layer = inputs_i
+            for ii in range(0,len(self.n_neuron)) :
+                if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
+                    layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec)
+                else :
+                    layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec)
+                if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+            # (nframes x natoms) x naxis
+            final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec, final_layer = True)
+            if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+            # (nframes x natoms) x 1 * naxis
+            final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], 1, self.dim_rot_mat_1])
+            # (nframes x natoms) x 1 x 3(coord)
+            final_layer = tf.matmul(final_layer, rot_mat_i)
+            # nframes x natoms x 3
+            final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3])
+
+            # concat the results
+            outs_list.append(final_layer)
+            count += 1
+        outs = tf.concat(outs_list, axis = 1)
+
+        tf.summary.histogram('fitting_net_output', outs)
+        return tf.reshape(outs, [-1])
+        # return tf.reshape(outs, [tf.shape(inputs)[0] * natoms[0] * 3 // 3])
+
+    def init_variables(self,
+                       graph: tf.Graph,
+                       graph_def: tf.GraphDef,
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the fitting net variables with the given dict
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        suffix : str
+            suffix to name scope
+        """
+        self.fitting_net_variables = get_fitting_net_variables_from_graph_def(graph_def, suffix=suffix)
+
+
+    def enable_mixed_precision(self, mixed_prec : dict = None) -> None:
+        """
+        Reveive the mixed precision setting.
+
+        Parameters
+        ----------
+        mixed_prec
+                The mixed precision setting used in the embedding net
+        """
+        self.mixed_prec = mixed_prec
+        self.fitting_precision = get_precision(mixed_prec['output_prec'])

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/fit/ener.py

+import warnings
+import numpy as np
+from typing import Tuple, List
+from packaging.version import Version
+
+from deepmd.env import tf
+from deepmd.common import add_data_requirement, get_activation_func, get_precision, cast_precision
+from deepmd.utils.network import one_layer_rand_seed_shift
+from deepmd.utils.network import one_layer as one_layer_deepmd
+from deepmd.utils.type_embed import embed_atom_type
+from deepmd.utils.graph import get_fitting_net_variables_from_graph_def, load_graph_def, get_tensor_by_name_from_graph
+from deepmd.utils.errors import GraphWithoutTensorError
+from deepmd.fit.fitting import Fitting
+
+from deepmd.env import global_cvt_2_tf_float
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION, TF_VERSION
+
+from deepmd.nvnmd.utils.config import nvnmd_cfg
+from deepmd.nvnmd.fit.ener import one_layer_nvnmd
+
+class EnerFitting (Fitting):
+    r"""Fitting the energy of the system. The force and the virial can also be trained.
+
+    The potential energy :math:`E` is a fitting network function of the descriptor :math:`\mathcal{D}`:
+
+    .. math::
+        E(\mathcal{D}) = \mathcal{L}^{(n)} \circ \mathcal{L}^{(n-1)}
+        \circ \cdots \circ \mathcal{L}^{(1)} \circ \mathcal{L}^{(0)}
+
+    The first :math:`n` hidden layers :math:`\mathcal{L}^{(0)}, \cdots, \mathcal{L}^{(n-1)}` are given by
+
+    .. math::
+        \mathbf{y}=\mathcal{L}(\mathbf{x};\mathbf{w},\mathbf{b})=
+            \boldsymbol{\phi}(\mathbf{x}^T\mathbf{w}+\mathbf{b})
+
+    where :math:`\mathbf{x} \in \mathbb{R}^{N_1}`$` is the input vector and :math:`\mathbf{y} \in \mathbb{R}^{N_2}`
+    is the output vector. :math:`\mathbf{w} \in \mathbb{R}^{N_1 \times N_2}` and
+    :math:`\mathbf{b} \in \mathbb{R}^{N_2}`$` are weights and biases, respectively,
+    both of which are trainable if `trainable[i]` is `True`. :math:`\boldsymbol{\phi}`
+    is the activation function.
+
+    The output layer :math:`\mathcal{L}^{(n)}` is given by
+
+    .. math::
+        \mathbf{y}=\mathcal{L}^{(n)}(\mathbf{x};\mathbf{w},\mathbf{b})=
+            \mathbf{x}^T\mathbf{w}+\mathbf{b}
+
+    where :math:`\mathbf{x} \in \mathbb{R}^{N_{n-1}}`$` is the input vector and :math:`\mathbf{y} \in \mathbb{R}`
+    is the output scalar. :math:`\mathbf{w} \in \mathbb{R}^{N_{n-1}}` and
+    :math:`\mathbf{b} \in \mathbb{R}`$` are weights and bias, respectively,
+    both of which are trainable if `trainable[n]` is `True`.
+
+    Parameters
+    ----------
+    descrpt
+            The descrptor :math:`\mathcal{D}`
+    neuron
+            Number of neurons :math:`N` in each hidden layer of the fitting net
+    resnet_dt
+            Time-step `dt` in the resnet construction:
+            :math:`y = x + dt * \phi (Wx + b)`
+    numb_fparam
+            Number of frame parameter
+    numb_aparam
+            Number of atomic parameter
+    rcond
+            The condition number for the regression of atomic energy.
+    tot_ener_zero
+            Force the total energy to zero. Useful for the charge fitting.
+    trainable
+            If the weights of fitting net are trainable. 
+            Suppose that we have :math:`N_l` hidden layers in the fitting net, 
+            this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable.
+    seed
+            Random seed for initializing the network parameters.
+    atom_ener
+            Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set.
+    activation_function
+            The activation function :math:`\boldsymbol{\phi}` in the embedding net. Supported options are |ACTIVATION_FN|
+    precision
+            The precision of the embedding net parameters. Supported options are |PRECISION|
+    uniform_seed
+            Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
+    """
+    def __init__ (self, 
+                  descrpt : tf.Tensor,
+                  neuron : List[int] = [120,120,120],
+                  resnet_dt : bool = True,
+                  numb_fparam : int = 0,
+                  numb_aparam : int = 0,
+                  rcond : float = 1e-3,
+                  tot_ener_zero : bool = False,
+                  trainable : List[bool] = None,
+                  seed : int = None,
+                  atom_ener : List[float] = [],
+                  activation_function : str = 'tanh',
+                  precision : str = 'default',
+                  uniform_seed: bool = False
+    ) -> None:
+        """
+        Constructor
+        """
+        # model param
+        self.ntypes = descrpt.get_ntypes()
+        self.dim_descrpt = descrpt.get_dim_out()
+        # args = ()\
+        #        .add('numb_fparam',      int,    default = 0)\
+        #        .add('numb_aparam',      int,    default = 0)\
+        #        .add('neuron',           list,   default = [120,120,120], alias = 'n_neuron')\
+        #        .add('resnet_dt',        bool,   default = True)\
+        #        .add('rcond',            float,  default = 1e-3) \
+        #        .add('tot_ener_zero',    bool,   default = False) \
+        #        .add('seed',             int)               \
+        #        .add('atom_ener',        list,   default = [])\
+        #        .add("activation_function", str,    default = "tanh")\
+        #        .add("precision",           str, default = "default")\
+        #        .add("trainable",        [list, bool], default = True)
+        self.numb_fparam = numb_fparam
+        self.numb_aparam = numb_aparam
+        self.n_neuron = neuron
+        self.resnet_dt = resnet_dt
+        self.rcond = rcond
+        self.seed = seed
+        self.uniform_seed = uniform_seed
+        self.seed_shift = one_layer_rand_seed_shift()
+        self.tot_ener_zero = tot_ener_zero
+        self.fitting_activation_fn = get_activation_func(activation_function)
+        self.fitting_precision = get_precision(precision)
+        self.trainable = trainable
+        if self.trainable is None:
+            self.trainable = [True for ii in range(len(self.n_neuron) + 1)]
+        if type(self.trainable) is bool:
+            self.trainable = [self.trainable] * (len(self.n_neuron)+1)
+        assert(len(self.trainable) == len(self.n_neuron) + 1), 'length of trainable should be that of n_neuron + 1'
+        self.atom_ener = []
+        self.atom_ener_v = atom_ener
+        for at, ae in enumerate(atom_ener):
+            if ae is not None:
+                self.atom_ener.append(tf.constant(ae, self.fitting_precision, name = "atom_%d_ener" % at))
+            else:
+                self.atom_ener.append(None)
+        self.useBN = False
+        self.bias_atom_e = np.zeros(self.ntypes, dtype=np.float64)
+        # data requirement
+        if self.numb_fparam > 0 :
+            add_data_requirement('fparam', self.numb_fparam, atomic=False, must=True, high_prec=False)
+            self.fparam_avg = None
+            self.fparam_std = None
+            self.fparam_inv_std = None
+        if self.numb_aparam > 0:
+            add_data_requirement('aparam', self.numb_aparam, atomic=True,  must=True, high_prec=False)
+            self.aparam_avg = None
+            self.aparam_std = None
+            self.aparam_inv_std = None
+
+        self.fitting_net_variables = None
+        self.mixed_prec = None
+
+    def get_numb_fparam(self) -> int:
+        """
+        Get the number of frame parameters
+        """
+        return self.numb_fparam
+
+    def get_numb_aparam(self) -> int:
+        """
+        Get the number of atomic parameters
+        """
+        return self.numb_fparam
+
+    def compute_output_stats(self, 
+                             all_stat: dict,
+                             mixed_type: bool = False
+    ) -> None:
+        """
+        Compute the ouput statistics
+
+        Parameters
+        ----------
+        all_stat
+                must have the following components:
+                all_stat['energy'] of shape n_sys x n_batch x n_frame
+                can be prepared by model.make_stat_input
+        mixed_type
+                Whether to perform the mixed_type mode.
+                If True, the input data has the mixed_type format (see doc/model/train_se_atten.md),
+                in which frames in a system may have different natoms_vec(s), with the same nloc.
+        """
+        self.bias_atom_e = self._compute_output_stats(all_stat, rcond=self.rcond, mixed_type=mixed_type)
+
+    def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False):
+        data = all_stat['energy']
+        # data[sys_idx][batch_idx][frame_idx]
+        sys_ener = np.array([])
+        for ss in range(len(data)):
+            sys_data = []
+            for ii in range(len(data[ss])):
+                for jj in range(len(data[ss][ii])):
+                    sys_data.append(data[ss][ii][jj])
+            sys_data = np.concatenate(sys_data)
+            sys_ener = np.append(sys_ener, np.average(sys_data))
+        sys_tynatom = np.array([])
+        if mixed_type:
+            data = all_stat['real_natoms_vec']
+            nsys = len(data)
+            for ss in range(len(data)):
+                tmp_tynatom = []
+                for ii in range(len(data[ss])):
+                    for jj in range(len(data[ss][ii])):
+                        tmp_tynatom.append(data[ss][ii][jj].astype(np.float64))
+                tmp_tynatom = np.average(np.array(tmp_tynatom), axis=0)
+                sys_tynatom = np.append(sys_tynatom, tmp_tynatom)
+        else:
+            data = all_stat['natoms_vec']
+            nsys = len(data)
+            for ss in range(len(data)):
+                sys_tynatom = np.append(sys_tynatom, data[ss][0].astype(np.float64))
+        sys_tynatom = np.reshape(sys_tynatom, [nsys,-1])
+        sys_tynatom = sys_tynatom[:,2:]
+        if len(self.atom_ener) > 0:
+            # Atomic energies stats are incorrect if atomic energies are assigned.
+            # In this situation, we directly use these assigned energies instead of computing stats.
+            # This will make the loss decrease quickly
+            assigned_atom_ener = np.array(list((ee for ee in self.atom_ener_v if ee is not None)))
+            assigned_ener_idx = list((ii for ii, ee in enumerate(self.atom_ener_v) if ee is not None))
+            # np.dot out size: nframe
+            sys_ener -= np.dot(sys_tynatom[:, assigned_ener_idx], assigned_atom_ener)
+            sys_tynatom[:, assigned_ener_idx] = 0.
+        energy_shift,resd,rank,s_value \
+            = np.linalg.lstsq(sys_tynatom, sys_ener, rcond = rcond)
+        if len(self.atom_ener) > 0:
+            for ii in assigned_ener_idx:
+                energy_shift[ii] = self.atom_ener_v[ii]
+        return energy_shift    
+
+    def compute_input_stats(self, 
+                            all_stat : dict,
+                            protection : float = 1e-2) -> None:
+        """
+        Compute the input statistics
+
+        Parameters
+        ----------
+        all_stat
+                if numb_fparam > 0 must have all_stat['fparam']
+                if numb_aparam > 0 must have all_stat['aparam']
+                can be prepared by model.make_stat_input
+        protection
+                Divided-by-zero protection
+        """
+        # stat fparam
+        if self.numb_fparam > 0:
+            cat_data = np.concatenate(all_stat['fparam'], axis = 0)
+            cat_data = np.reshape(cat_data, [-1, self.numb_fparam])
+            self.fparam_avg = np.average(cat_data, axis = 0)
+            self.fparam_std = np.std(cat_data, axis = 0)
+            for ii in range(self.fparam_std.size):
+                if self.fparam_std[ii] < protection:
+                    self.fparam_std[ii] = protection
+            self.fparam_inv_std = 1./self.fparam_std
+        # stat aparam
+        if self.numb_aparam > 0:
+            sys_sumv = []
+            sys_sumv2 = []
+            sys_sumn = []
+            for ss_ in all_stat['aparam'] : 
+                ss = np.reshape(ss_, [-1, self.numb_aparam])
+                sys_sumv.append(np.sum(ss, axis = 0))
+                sys_sumv2.append(np.sum(np.multiply(ss, ss), axis = 0))
+                sys_sumn.append(ss.shape[0])
+            sumv = np.sum(sys_sumv, axis = 0)
+            sumv2 = np.sum(sys_sumv2, axis = 0)
+            sumn = np.sum(sys_sumn)
+            self.aparam_avg = (sumv)/sumn
+            self.aparam_std = self._compute_std(sumv2, sumv, sumn)
+            for ii in range(self.aparam_std.size):
+                if self.aparam_std[ii] < protection:
+                    self.aparam_std[ii] = protection
+            self.aparam_inv_std = 1./self.aparam_std
+
+
+    def _compute_std (self, sumv2, sumv, sumn) :
+        return np.sqrt(sumv2/sumn - np.multiply(sumv/sumn, sumv/sumn))
+
+    def _build_lower(
+            self,
+            start_index,
+            natoms,
+            inputs,
+            fparam = None,
+            aparam = None, 
+            bias_atom_e = 0.0,
+            suffix = '',
+            reuse = None
+    ):
+        # cut-out inputs
+        inputs_i = tf.slice (inputs,
+                             [ 0, start_index, 0],
+                             [-1, natoms, -1] )
+        inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
+        layer = inputs_i
+        if fparam is not None:
+            ext_fparam = tf.tile(fparam, [1, natoms])
+            ext_fparam = tf.reshape(ext_fparam, [-1, self.numb_fparam])
+            ext_fparam = tf.cast(ext_fparam,self.fitting_precision)
+            layer = tf.concat([layer, ext_fparam], axis = 1)
+        if aparam is not None:
+            ext_aparam = tf.slice(aparam, 
+                                  [ 0, start_index      * self.numb_aparam],
+                                  [-1, natoms * self.numb_aparam])
+            ext_aparam = tf.reshape(ext_aparam, [-1, self.numb_aparam])
+            ext_aparam = tf.cast(ext_aparam,self.fitting_precision)
+            layer = tf.concat([layer, ext_aparam], axis = 1)
+
+        if nvnmd_cfg.enable: 
+            one_layer = one_layer_nvnmd
+        else:
+            one_layer = one_layer_deepmd
+        for ii in range(0,len(self.n_neuron)) :
+            if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] and (not nvnmd_cfg.enable):
+                layer+= one_layer(
+                    layer,
+                    self.n_neuron[ii],
+                    name='layer_'+str(ii)+suffix,
+                    reuse=reuse,
+                    seed = self.seed,
+                    use_timestep = self.resnet_dt,
+                    activation_fn = self.fitting_activation_fn,
+                    precision = self.fitting_precision,
+                    trainable = self.trainable[ii],
+                    uniform_seed = self.uniform_seed,
+                    initial_variables = self.fitting_net_variables,
+                    mixed_prec = self.mixed_prec)
+            else :
+                layer = one_layer(
+                    layer,
+                    self.n_neuron[ii],
+                    name='layer_'+str(ii)+suffix,
+                    reuse=reuse,
+                    seed = self.seed,
+                    activation_fn = self.fitting_activation_fn,
+                    precision = self.fitting_precision,
+                    trainable = self.trainable[ii],
+                    uniform_seed = self.uniform_seed,
+                    initial_variables = self.fitting_net_variables,
+                    mixed_prec = self.mixed_prec)
+            if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+        final_layer = one_layer(
+            layer, 
+            1, 
+            activation_fn = None, 
+            bavg = bias_atom_e, 
+            name='final_layer'+suffix, 
+            reuse=reuse, 
+            seed = self.seed, 
+            precision = self.fitting_precision, 
+            trainable = self.trainable[-1],
+            uniform_seed = self.uniform_seed,
+            initial_variables = self.fitting_net_variables,
+            mixed_prec = self.mixed_prec,
+            final_layer = True)
+        if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+
+        return final_layer
+            
+            
+    @cast_precision
+    def build (self, 
+               inputs : tf.Tensor,
+               natoms : tf.Tensor,
+               input_dict : dict = None,
+               reuse : bool = None,
+               suffix : str = '', 
+    ) -> tf.Tensor:
+        """
+        Build the computational graph for fitting net
+
+        Parameters
+        ----------
+        inputs
+                The input descriptor
+        input_dict
+                Additional dict for inputs. 
+                if numb_fparam > 0, should have input_dict['fparam']
+                if numb_aparam > 0, should have input_dict['aparam']
+        natoms
+                The number of atoms. This tensor has the length of Ntypes + 2
+                natoms[0]: number of local atoms
+                natoms[1]: total number of atoms held by this processor
+                natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
+        reuse
+                The weights in the networks should be reused when get the variable.
+        suffix
+                Name suffix to identify this descriptor
+
+        Returns
+        -------
+        ener
+                The system energy
+        """
+        if input_dict is None:
+            input_dict = {}
+        bias_atom_e = self.bias_atom_e
+        type_embedding = input_dict.get('type_embedding', None)
+        atype = input_dict.get('atype', None)
+        if self.numb_fparam > 0:
+            if self.fparam_avg is None:
+                self.fparam_avg = 0.
+            if self.fparam_inv_std is None:
+                self.fparam_inv_std = 1.
+        if self.numb_aparam > 0:
+            if self.aparam_avg is None:
+                self.aparam_avg = 0.
+            if self.aparam_inv_std is None:
+                self.aparam_inv_std = 1.
+
+        with tf.variable_scope('fitting_attr' + suffix, reuse = reuse) :
+            t_dfparam = tf.constant(self.numb_fparam, 
+                                    name = 'dfparam', 
+                                    dtype = tf.int32)
+            t_daparam = tf.constant(self.numb_aparam, 
+                                    name = 'daparam', 
+                                    dtype = tf.int32)
+            if type_embedding is not None:
+                self.t_bias_atom_e = tf.get_variable('t_bias_atom_e',
+                                            self.bias_atom_e.shape,
+                                            dtype=self.fitting_precision,
+                                            trainable=False,
+                                            initializer=tf.constant_initializer(self.bias_atom_e))
+            if self.numb_fparam > 0: 
+                t_fparam_avg = tf.get_variable('t_fparam_avg', 
+                                               self.numb_fparam,
+                                               dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                               trainable = False,
+                                               initializer = tf.constant_initializer(self.fparam_avg))
+                t_fparam_istd = tf.get_variable('t_fparam_istd', 
+                                                self.numb_fparam,
+                                                dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                                trainable = False,
+                                                initializer = tf.constant_initializer(self.fparam_inv_std))
+            if self.numb_aparam > 0: 
+                t_aparam_avg = tf.get_variable('t_aparam_avg', 
+                                               self.numb_aparam,
+                                               dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                               trainable = False,
+                                               initializer = tf.constant_initializer(self.aparam_avg))
+                t_aparam_istd = tf.get_variable('t_aparam_istd', 
+                                                self.numb_aparam,
+                                                dtype = GLOBAL_TF_FLOAT_PRECISION,
+                                                trainable = False,
+                                                initializer = tf.constant_initializer(self.aparam_inv_std))
+            
+        inputs = tf.reshape(inputs, [-1, natoms[0], self.dim_descrpt])
+        if len(self.atom_ener):
+            # only for atom_ener
+            nframes = input_dict.get('nframes')
+            if nframes is not None:
+                # like inputs, but we don't want to add a dependency on inputs
+                inputs_zero = tf.zeros((nframes, natoms[0], self.dim_descrpt), dtype=self.fitting_precision)
+            else:
+                inputs_zero = tf.zeros_like(inputs, dtype=self.fitting_precision)
+        
+
+        if bias_atom_e is not None :
+            assert(len(bias_atom_e) == self.ntypes)
+
+        fparam = None
+        aparam = None
+        if self.numb_fparam > 0 :
+            fparam = input_dict['fparam']
+            fparam = tf.reshape(fparam, [-1, self.numb_fparam])
+            fparam = (fparam - t_fparam_avg) * t_fparam_istd            
+        if self.numb_aparam > 0 :
+            aparam = input_dict['aparam']
+            aparam = tf.reshape(aparam, [-1, self.numb_aparam])
+            aparam = (aparam - t_aparam_avg) * t_aparam_istd
+            aparam = tf.reshape(aparam, [-1, self.numb_aparam * natoms[0]])
+
+        if type_embedding is not None:
+            atype_nall = tf.reshape(atype, [-1, natoms[1]])
+            self.atype_nloc = tf.reshape(tf.slice(atype_nall, [0, 0], [-1, natoms[0]]), [-1])  ## lammps will make error
+            atype_embed = tf.nn.embedding_lookup(type_embedding, self.atype_nloc)
+        else:
+            atype_embed = None
+
+        self.atype_embed = atype_embed
+
+        if atype_embed is None:
+            start_index = 0
+            outs_list = []
+            for type_i in range(self.ntypes):
+                if bias_atom_e is None :
+                    type_bias_ae = 0.0
+                else :
+                    type_bias_ae = bias_atom_e[type_i]
+                final_layer = self._build_lower(
+                    start_index, natoms[2+type_i], 
+                    inputs, fparam, aparam, 
+                    bias_atom_e=type_bias_ae, suffix='_type_'+str(type_i)+suffix, reuse=reuse
+                )
+                # concat the results
+                if type_i < len(self.atom_ener) and self.atom_ener[type_i] is not None:                
+                    zero_layer = self._build_lower(
+                        start_index, natoms[2+type_i], 
+                        inputs_zero, fparam, aparam, 
+                        bias_atom_e=type_bias_ae, suffix='_type_'+str(type_i)+suffix, reuse=True
+                    )
+                    final_layer += self.atom_ener[type_i] - zero_layer
+                final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i]])
+                outs_list.append(final_layer)
+                start_index += natoms[2+type_i]
+            # concat the results
+            # concat once may be faster than multiple concat
+            outs = tf.concat(outs_list, axis = 1)
+        # with type embedding
+        else:
+            if len(self.atom_ener) > 0:
+                raise RuntimeError("setting atom_ener is not supported by type embedding")
+            atype_embed = tf.cast(atype_embed, self.fitting_precision)
+            type_shape = atype_embed.get_shape().as_list()
+            inputs = tf.concat(
+                [tf.reshape(inputs,[-1,self.dim_descrpt]),atype_embed],
+                axis=1
+            )
+            self.dim_descrpt = self.dim_descrpt + type_shape[1]
+            inputs = tf.reshape(inputs, [-1, natoms[0], self.dim_descrpt])
+            final_layer = self._build_lower(
+                0, natoms[0], 
+                inputs, fparam, aparam, 
+                bias_atom_e=0.0, suffix=suffix, reuse=reuse
+            )
+            outs = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[0]])
+            # add bias
+            self.atom_ener_before = outs
+            self.add_type = tf.reshape(tf.nn.embedding_lookup(self.t_bias_atom_e, self.atype_nloc), [tf.shape(inputs)[0], natoms[0]])
+            outs = outs + self.add_type
+            self.atom_ener_after = outs
+
+        if self.tot_ener_zero:
+            force_tot_ener = 0.0
+            outs = tf.reshape(outs, [-1, natoms[0]])
+            outs_mean = tf.reshape(tf.reduce_mean(outs, axis = 1), [-1, 1])
+            outs_mean = outs_mean - tf.ones_like(outs_mean, dtype = GLOBAL_TF_FLOAT_PRECISION) * (force_tot_ener/global_cvt_2_tf_float(natoms[0]))
+            outs = outs - outs_mean
+            outs = tf.reshape(outs, [-1])
+
+        tf.summary.histogram('fitting_net_output', outs)
+        return tf.reshape(outs, [-1])
+
+
+    def init_variables(self,
+                       graph: tf.Graph,
+                       graph_def: tf.GraphDef,
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the fitting net variables with the given dict
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        suffix : str
+            suffix to name scope
+        """
+        self.fitting_net_variables = get_fitting_net_variables_from_graph_def(graph_def, suffix=suffix)
+        if self.numb_fparam > 0:
+            self.fparam_avg = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_fparam_avg' % suffix)
+            self.fparam_inv_std = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_fparam_istd' % suffix)
+        if self.numb_aparam > 0:
+            self.aparam_avg = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_aparam_avg' % suffix)
+            self.aparam_inv_std = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_aparam_istd' % suffix)
+        try:
+            self.bias_atom_e = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_bias_atom_e' % suffix)
+        except GraphWithoutTensorError:
+            # model without type_embedding has no t_bias_atom_e
+            pass
+
+    def enable_compression(self,
+                           model_file: str,
+                           suffix: str = ""
+    ) -> None:
+        """
+        Set the fitting net attributes from the frozen model_file when fparam or aparam is not zero
+
+        Parameters
+        ----------
+        model_file : str
+            The input frozen model file
+        suffix : str, optional
+                The suffix of the scope
+        """
+        if self.numb_fparam > 0 or self.numb_aparam > 0:
+            graph, _ = load_graph_def(model_file)
+        if self.numb_fparam > 0:
+            self.fparam_avg = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_fparam_avg' % suffix)
+            self.fparam_inv_std = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_fparam_istd' % suffix)
+        if self.numb_aparam > 0:
+            self.aparam_avg = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_aparam_avg' % suffix)
+            self.aparam_inv_std = get_tensor_by_name_from_graph(graph, 'fitting_attr%s/t_aparam_istd' % suffix)
+ 
+
+    def enable_mixed_precision(self, mixed_prec: dict = None) -> None:
+        """
+        Reveive the mixed precision setting.
+
+        Parameters
+        ----------
+        mixed_prec
+                The mixed precision setting used in the embedding net
+        """
+        self.mixed_prec = mixed_prec
+        self.fitting_precision = get_precision(mixed_prec['output_prec'])

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/fit/fitting.py

+from deepmd.env import tf
+from deepmd.utils import Plugin, PluginVariant
+
+class Fitting:
+    @property
+    def precision(self) -> tf.DType:
+        """Precision of fitting network."""
+        return self.fitting_precision
+
+    def init_variables(self,
+                       graph: tf.Graph,
+                       graph_def: tf.GraphDef,
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the fitting net variables with the given dict
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        suffix : str
+            suffix to name scope
+        
+        Notes
+        -----
+        This method is called by others when the fitting supported initialization from the given variables.
+        """
+        raise NotImplementedError(
+            "Fitting %s doesn't support initialization from the given variables!" % type(self).__name__)

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/fit/polar.py

+import warnings
+import numpy as np
+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.common import add_data_requirement, cast_precision, get_activation_func, get_precision
+from deepmd.utils.network import one_layer, one_layer_rand_seed_shift
+from deepmd.utils.graph import get_fitting_net_variables_from_graph_def
+from deepmd.descriptor import DescrptLocFrame
+from deepmd.descriptor import DescrptSeA
+from deepmd.fit.fitting import Fitting
+
+from deepmd.env import global_cvt_2_tf_float
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+
+
+class PolarFittingLocFrame () :
+    """
+    Fitting polarizability with local frame descriptor.
+
+    .. deprecated:: 2.0.0
+        This class is not supported any more.
+    """
+    def __init__ (self, jdata, descrpt) :
+        if not isinstance(descrpt, DescrptLocFrame) :
+            raise RuntimeError('PolarFittingLocFrame only supports DescrptLocFrame')
+        self.ntypes = descrpt.get_ntypes()
+        self.dim_descrpt = descrpt.get_dim_out()
+        args = ClassArg()\
+               .add('neuron',           list, default = [120,120,120], alias = 'n_neuron')\
+               .add('resnet_dt',        bool, default = True)\
+               .add('sel_type',         [list,int], default = [ii for ii in range(self.ntypes)], alias = 'pol_type')\
+               .add('seed',             int)\
+               .add("activation_function", str, default = "tanh")\
+               .add('precision',           str,    default = "default")    
+        class_data = args.parse(jdata)
+        self.n_neuron = class_data['neuron']
+        self.resnet_dt = class_data['resnet_dt']
+        self.sel_type = class_data['sel_type']
+        self.seed = class_data['seed']
+        self.fitting_activation_fn = get_activation_func(class_data["activation_function"])
+        self.fitting_precision = get_precision(class_data['precision'])
+        self.useBN = False
+
+    def get_sel_type(self):
+        return self.sel_type
+
+    def get_out_size(self):
+        return 9
+
+    def build (self, 
+               input_d,
+               rot_mat,
+               natoms,
+               reuse = None,
+               suffix = '') :
+        start_index = 0
+        inputs = tf.cast(tf.reshape(input_d, [-1, natoms[0], self.dim_descrpt]), self.fitting_precision)
+        rot_mat = tf.reshape(rot_mat, [-1, 9 * natoms[0]])
+
+        count = 0
+        outs_list = []
+        for type_i in range(self.ntypes):
+            # cut-out inputs
+            inputs_i = tf.slice (inputs,
+                                 [ 0, start_index, 0],
+                                 [-1, natoms[2+type_i], -1] )
+            inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
+            rot_mat_i = tf.slice (rot_mat,
+                                  [ 0, start_index*      9],
+                                  [-1, natoms[2+type_i]* 9] )
+            rot_mat_i = tf.reshape(rot_mat_i, [-1, 3, 3])
+            start_index += natoms[2+type_i]
+            if not type_i in self.sel_type :
+                continue
+            layer = inputs_i
+            for ii in range(0,len(self.n_neuron)) :
+                if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
+                    layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision)
+                else :
+                    layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision)
+            # (nframes x natoms) x 9
+            final_layer = one_layer(layer, 9, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, final_layer = True)
+            # (nframes x natoms) x 3 x 3
+            final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], 3, 3])
+            # (nframes x natoms) x 3 x 3
+            final_layer = final_layer + tf.transpose(final_layer, perm = [0,2,1])
+            # (nframes x natoms) x 3 x 3(coord)
+            final_layer = tf.matmul(final_layer, rot_mat_i)
+            # (nframes x natoms) x 3(coord) x 3(coord)
+            final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a = True)
+            # nframes x natoms x 3 x 3
+            final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3, 3])
+
+            # concat the results
+            outs_list.append(final_layer)
+            count += 1
+        outs = tf.concat(outs_list, axis = 1)
+
+        tf.summary.histogram('fitting_net_output', outs)
+        return tf.cast(tf.reshape(outs, [-1]),  GLOBAL_TF_FLOAT_PRECISION)
+
+
+class PolarFittingSeA (Fitting) :
+    """
+    Fit the atomic polarizability with descriptor se_a
+
+    Parameters
+    ----------
+    descrpt : tf.Tensor
+            The descrptor
+    neuron : List[int]
+            Number of neurons in each hidden layer of the fitting net
+    resnet_dt : bool
+            Time-step `dt` in the resnet construction:
+            y = x + dt * \phi (Wx + b)
+    sel_type : List[int]
+            The atom types selected to have an atomic polarizability prediction. If is None, all atoms are selected.
+    fit_diag : bool
+            Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix.
+    scale : List[float]
+            The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i]
+    diag_shift : List[float]
+            The diagonal part of the polarizability matrix of type i will be shifted by diag_shift[i]. The shift operation is carried out after scale.        
+    seed : int
+            Random seed for initializing the network parameters.
+    activation_function : str
+            The activation function in the embedding net. Supported options are |ACTIVATION_FN|
+    precision : str
+            The precision of the embedding net parameters. Supported options are |PRECISION|
+    uniform_seed
+            Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
+    """
+    def __init__ (self, 
+                  descrpt : tf.Tensor,
+                  neuron : List[int] = [120,120,120],
+                  resnet_dt : bool = True,
+                  sel_type : List[int] = None,
+                  fit_diag : bool = True,
+                  scale : List[float] = None,
+                  shift_diag : bool = True,     # YWolfeee: will support the user to decide whether to use this function
+                  #diag_shift : List[float] = None, YWolfeee: will not support the user to assign a shift
+                  seed : int = None,
+                  activation_function : str = 'tanh',
+                  precision : str = 'default',
+                  uniform_seed: bool = False                  
+    ) -> None:
+        """
+        Constructor
+        """
+        if not isinstance(descrpt, DescrptSeA) :
+            raise RuntimeError('PolarFittingSeA only supports DescrptSeA')
+        self.ntypes = descrpt.get_ntypes()
+        self.dim_descrpt = descrpt.get_dim_out()
+        # args = ClassArg()\
+        #        .add('neuron',           list,   default = [120,120,120], alias = 'n_neuron')\
+        #        .add('resnet_dt',        bool,   default = True)\
+        #        .add('fit_diag',         bool,   default = True)\
+        #        .add('diag_shift',       [list,float], default = [0.0 for ii in range(self.ntypes)])\
+        #        .add('scale',            [list,float], default = [1.0 for ii in range(self.ntypes)])\
+        #        .add('sel_type',         [list,int],   default = [ii for ii in range(self.ntypes)], alias = 'pol_type')\
+        #        .add('seed',             int)\
+        #        .add("activation_function", str ,   default = "tanh")\
+        #        .add('precision',           str,    default = "default")
+        # class_data = args.parse(jdata)
+        self.n_neuron = neuron
+        self.resnet_dt = resnet_dt
+        self.sel_type = sel_type
+        self.fit_diag = fit_diag
+        self.seed = seed
+        self.uniform_seed = uniform_seed
+        self.seed_shift = one_layer_rand_seed_shift()
+        #self.diag_shift = diag_shift
+        self.shift_diag = shift_diag
+        self.scale = scale
+        self.fitting_activation_fn = get_activation_func(activation_function)
+        self.fitting_precision = get_precision(precision)
+        if self.sel_type is None:
+            self.sel_type = [ii for ii in range(self.ntypes)]
+        if self.scale is None:
+            self.scale = [1.0 for ii in range(self.ntypes)]
+        #if self.diag_shift is None:
+        #    self.diag_shift = [0.0 for ii in range(self.ntypes)]
+        if type(self.sel_type) is not list:
+            self.sel_type = [self.sel_type]
+        self.constant_matrix = np.zeros(len(self.sel_type)) # len(sel_type) x 1, store the average diagonal value
+        #if type(self.diag_shift) is not list:
+        #    self.diag_shift = [self.diag_shift]
+        if type(self.scale) is not list:
+            self.scale = [self.scale]
+        self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1()
+        self.dim_rot_mat = self.dim_rot_mat_1 * 3
+        self.useBN = False
+        self.fitting_net_variables = None
+        self.mixed_prec = None
+
+    def get_sel_type(self) -> List[int]:
+        """
+        Get selected atom types
+        """
+        return self.sel_type
+
+    def get_out_size(self) -> int:
+        """
+        Get the output size. Should be 9
+        """
+        return 9
+
+    def compute_input_stats(self, 
+                            all_stat, 
+                            protection = 1e-2):
+        """
+        Compute the input statistics
+
+        Parameters
+        ----------
+        all_stat
+                Dictionary of inputs. 
+                can be prepared by model.make_stat_input
+        protection
+                Divided-by-zero protection
+        """
+        if not ('polarizability' in all_stat.keys()):
+            self.avgeig = np.zeros([9])
+            warnings.warn('no polarizability data, cannot do data stat. use zeros as guess')
+            return
+        data = all_stat['polarizability']
+        all_tmp = []
+        for ss in range(len(data)):
+            tmp = np.concatenate(data[ss], axis = 0)
+            tmp = np.reshape(tmp, [-1, 3, 3])
+            tmp,_ = np.linalg.eig(tmp)
+            tmp = np.absolute(tmp)
+            tmp = np.sort(tmp, axis = 1)
+            all_tmp.append(tmp)
+        all_tmp = np.concatenate(all_tmp, axis = 1)
+        self.avgeig = np.average(all_tmp, axis = 0)
+
+        # YWolfeee: support polar normalization, initialize to a more appropriate point 
+        if self.shift_diag:
+            mean_polar = np.zeros([len(self.sel_type), 9])
+            sys_matrix, polar_bias = [], []
+            for ss in range(len(all_stat['type'])):
+                atom_has_polar = [w for w in all_stat['type'][ss][0] if (w in self.sel_type)]   # select atom with polar
+                if all_stat['find_atomic_polarizability'][ss] > 0.0:
+                    for itype in range(len(self.sel_type)): # Atomic polar mode, should specify the atoms
+                        index_lis = [index for index, w in enumerate(atom_has_polar) \
+                                        if atom_has_polar[index] == self.sel_type[itype]]   # select index in this type
+
+                        sys_matrix.append(np.zeros((1,len(self.sel_type))))
+                        sys_matrix[-1][0,itype] = len(index_lis)
+
+                        polar_bias.append(np.sum(
+                            all_stat['atomic_polarizability'][ss].reshape((-1,9))[index_lis],axis=0).reshape((1,9)))
+                else:   # No atomic polar in this system, so it should have global polar
+                    if not all_stat['find_polarizability'][ss] > 0.0: # This system is jsut a joke?
+                        continue
+                    # Till here, we have global polar
+                    sys_matrix.append(np.zeros((1,len(self.sel_type)))) # add a line in the equations
+                    for itype in range(len(self.sel_type)): # Atomic polar mode, should specify the atoms
+                        index_lis = [index for index, w in enumerate(atom_has_polar) \
+                                        if atom_has_polar[index] == self.sel_type[itype]]   # select index in this type
+
+                        sys_matrix[-1][0,itype] = len(index_lis)
+                    
+                    # add polar_bias
+                    polar_bias.append(all_stat['polarizability'][ss].reshape((1,9)))
+
+            matrix, bias = np.concatenate(sys_matrix,axis=0), np.concatenate(polar_bias,axis=0)
+            atom_polar,_,_,_ \
+                = np.linalg.lstsq(matrix, bias, rcond = 1e-3)
+            for itype in range(len(self.sel_type)):
+                self.constant_matrix[itype] = np.mean(np.diagonal(atom_polar[itype].reshape((3,3))))
+
+    @cast_precision
+    def build (self, 
+               input_d : tf.Tensor,
+               rot_mat : tf.Tensor,
+               natoms : tf.Tensor,
+               reuse : bool = None,
+               suffix : str = '') :
+        """
+        Build the computational graph for fitting net
+        
+        Parameters
+        ----------
+        input_d
+                The input descriptor
+        rot_mat
+                The rotation matrix from the descriptor.
+        natoms
+                The number of atoms. This tensor has the length of Ntypes + 2
+                natoms[0]: number of local atoms
+                natoms[1]: total number of atoms held by this processor
+                natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
+        reuse
+                The weights in the networks should be reused when get the variable.
+        suffix
+                Name suffix to identify this descriptor
+
+        Returns
+        -------
+        atomic_polar
+                The atomic polarizability        
+        """
+        start_index = 0
+        inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]])
+        rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]])
+
+        count = 0
+        outs_list = []
+        for type_i in range(self.ntypes):
+            # cut-out inputs
+            inputs_i = tf.slice (inputs,
+                                 [ 0, start_index*      self.dim_descrpt],
+                                 [-1, natoms[2+type_i]* self.dim_descrpt] )
+            inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
+            rot_mat_i = tf.slice (rot_mat,
+                                  [ 0, start_index*      self.dim_rot_mat],
+                                  [-1, natoms[2+type_i]* self.dim_rot_mat] )
+            rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3])
+            start_index += natoms[2+type_i]
+            if not type_i in self.sel_type :
+                continue
+            layer = inputs_i
+            for ii in range(0,len(self.n_neuron)) :
+                if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
+                    layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec)
+                else :
+                    layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec)
+                if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+            if self.fit_diag :
+                bavg = np.zeros(self.dim_rot_mat_1)
+                # bavg[0] = self.avgeig[0]
+                # bavg[1] = self.avgeig[1]
+                # bavg[2] = self.avgeig[2]
+                # (nframes x natoms) x naxis
+                final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec, final_layer = True)
+                if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+                # (nframes x natoms) x naxis
+                final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.dim_rot_mat_1])
+                # (nframes x natoms) x naxis x naxis
+                final_layer = tf.matrix_diag(final_layer)                
+            else :
+                bavg = np.zeros(self.dim_rot_mat_1*self.dim_rot_mat_1)
+                # bavg[0*self.dim_rot_mat_1+0] = self.avgeig[0]
+                # bavg[1*self.dim_rot_mat_1+1] = self.avgeig[1]
+                # bavg[2*self.dim_rot_mat_1+2] = self.avgeig[2]
+                # (nframes x natoms) x (naxis x naxis)
+                final_layer = one_layer(layer, self.dim_rot_mat_1*self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed, initial_variables = self.fitting_net_variables, mixed_prec = self.mixed_prec, final_layer = True)
+                if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+                # (nframes x natoms) x naxis x naxis
+                final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.dim_rot_mat_1, self.dim_rot_mat_1])
+                # (nframes x natoms) x naxis x naxis
+                final_layer = final_layer + tf.transpose(final_layer, perm = [0,2,1])
+            # (nframes x natoms) x naxis x 3(coord)
+            final_layer = tf.matmul(final_layer, rot_mat_i)
+            # (nframes x natoms) x 3(coord) x 3(coord)
+            final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a = True)
+            # nframes x natoms x 3 x 3
+            final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3, 3])
+            # shift and scale
+            sel_type_idx = self.sel_type.index(type_i)
+            final_layer = final_layer * self.scale[sel_type_idx]
+            final_layer = final_layer + self.constant_matrix[sel_type_idx] * tf.eye(3, batch_shape=[tf.shape(inputs)[0], natoms[2+type_i]], dtype = GLOBAL_TF_FLOAT_PRECISION)
+
+            # concat the results
+            outs_list.append(final_layer)
+            count += 1
+        outs = tf.concat(outs_list, axis = 1)
+        
+        tf.summary.histogram('fitting_net_output', outs)
+        return tf.reshape(outs, [-1])
+
+    def init_variables(self,
+                       graph: tf.Graph,
+                       graph_def: tf.GraphDef,
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the fitting net variables with the given dict
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        suffix : str
+            suffix to name scope
+        """
+        self.fitting_net_variables = get_fitting_net_variables_from_graph_def(graph_def, suffix=suffix)
+
+
+    def enable_mixed_precision(self, mixed_prec : dict = None) -> None:
+        """
+        Reveive the mixed precision setting.
+
+        Parameters
+        ----------
+        mixed_prec
+                The mixed precision setting used in the embedding net
+        """
+        self.mixed_prec = mixed_prec
+        self.fitting_precision = get_precision(mixed_prec['output_prec'])
+
+
+class GlobalPolarFittingSeA () :
+    """
+    Fit the system polarizability with descriptor se_a
+
+    Parameters
+    ----------
+    descrpt : tf.Tensor
+            The descrptor
+    neuron : List[int]
+            Number of neurons in each hidden layer of the fitting net
+    resnet_dt : bool
+            Time-step `dt` in the resnet construction:
+            y = x + dt * \phi (Wx + b)
+    sel_type : List[int]
+            The atom types selected to have an atomic polarizability prediction
+    fit_diag : bool
+            Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix.
+    scale : List[float]
+            The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i]
+    diag_shift : List[float]
+            The diagonal part of the polarizability matrix of type i will be shifted by diag_shift[i]. The shift operation is carried out after scale.        
+    seed : int
+            Random seed for initializing the network parameters.
+    activation_function : str
+            The activation function in the embedding net. Supported options are |ACTIVATION_FN|
+    precision : str
+            The precision of the embedding net parameters. Supported options are |PRECISION|
+    """
+    def __init__ (self, 
+                  descrpt : tf.Tensor,
+                  neuron : List[int] = [120,120,120],
+                  resnet_dt : bool = True,
+                  sel_type : List[int] = None,
+                  fit_diag : bool = True,
+                  scale : List[float] = None,
+                  diag_shift : List[float] = None,
+                  seed : int = None,
+                  activation_function : str = 'tanh',
+                  precision : str = 'default'
+    ) -> None:
+        """
+        Constructor            
+        """
+        if not isinstance(descrpt, DescrptSeA) :
+            raise RuntimeError('GlobalPolarFittingSeA only supports DescrptSeA')
+        self.ntypes = descrpt.get_ntypes()
+        self.dim_descrpt = descrpt.get_dim_out()
+        self.polar_fitting = PolarFittingSeA(descrpt,
+                                             neuron,
+                                             resnet_dt,
+                                             sel_type,
+                                             fit_diag,
+                                             scale,
+                                             diag_shift,
+                                             seed,
+                                             activation_function,
+                                             precision)
+
+    def get_sel_type(self) -> int:
+        """
+        Get selected atom types
+        """
+        return self.polar_fitting.get_sel_type()
+
+    def get_out_size(self) -> int:
+        """
+        Get the output size. Should be 9
+        """
+        return self.polar_fitting.get_out_size()
+
+    def build (self,
+               input_d,
+               rot_mat,
+               natoms,
+               reuse = None,
+               suffix = '') -> tf.Tensor:
+        """
+        Build the computational graph for fitting net
+        
+        Parameters
+        ----------
+        input_d
+                The input descriptor
+        rot_mat
+                The rotation matrix from the descriptor.
+        natoms
+                The number of atoms. This tensor has the length of Ntypes + 2
+                natoms[0]: number of local atoms
+                natoms[1]: total number of atoms held by this processor
+                natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
+        reuse
+                The weights in the networks should be reused when get the variable.
+        suffix
+                Name suffix to identify this descriptor
+
+        Returns
+        -------
+        polar
+                The system polarizability        
+        """
+        inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]])
+        outs = self.polar_fitting.build(input_d, rot_mat, natoms, reuse, suffix)
+        # nframes x natoms x 9
+        outs = tf.reshape(outs, [tf.shape(inputs)[0], -1, 9])
+        outs = tf.reduce_sum(outs, axis = 1)
+        tf.summary.histogram('fitting_net_output', outs)
+        return tf.reshape(outs, [-1])
+    
+    def init_variables(self,
+                       graph: tf.Graph,
+                       graph_def: tf.GraphDef,
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the fitting net variables with the given dict
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        suffix : str
+            suffix to name scope
+        """
+        self.polar_fitting.init_variables(graph=graph, graph_def=graph_def, suffix=suffix)
+
+
+    def enable_mixed_precision(self, mixed_prec : dict = None) -> None:
+        """
+        Reveive the mixed precision setting.
+
+        Parameters
+        ----------
+        mixed_prec
+                The mixed precision setting used in the embedding net
+        """
+        self.polar_fitting.enable_mixed_precision(mixed_prec)

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/fit/wfc.py

+import warnings
+import numpy as np
+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.common import ClassArg, add_data_requirement, get_activation_func, get_precision
+from deepmd.utils.network import one_layer, one_layer_rand_seed_shift
+from deepmd.descriptor import DescrptLocFrame
+
+from deepmd.env import global_cvt_2_tf_float
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+
+class WFCFitting () :
+    """
+    Fitting Wannier function centers (WFCs) with local frame descriptor.
+
+    .. deprecated:: 2.0.0
+        This class is not supported any more.
+    """
+    def __init__ (self, jdata, descrpt):
+        if not isinstance(descrpt, DescrptLocFrame) :
+            raise RuntimeError('WFC only supports DescrptLocFrame')
+        self.ntypes = descrpt.get_ntypes()
+        self.dim_descrpt = descrpt.get_dim_out()
+        args = ClassArg()\
+               .add('neuron',           list,   default = [120,120,120], alias = 'n_neuron')\
+               .add('resnet_dt',        bool,   default = True)\
+               .add('wfc_numb',         int,    must = True)\
+               .add('sel_type',         [list,int],   default = [ii for ii in range(self.ntypes)], alias = 'wfc_type')\
+               .add('seed',             int)\
+               .add("activation_function", str, default = "tanh")\
+               .add('precision',           str,    default = "default")\
+               .add('uniform_seed',     bool, default = False)
+        class_data = args.parse(jdata)
+        self.n_neuron = class_data['neuron']
+        self.resnet_dt = class_data['resnet_dt']
+        self.wfc_numb = class_data['wfc_numb']
+        self.sel_type = class_data['sel_type']
+        self.seed = class_data['seed']
+        self.uniform_seed = class_data['uniform_seed']
+        self.seed_shift = one_layer_rand_seed_shift()
+        self.fitting_activation_fn = get_activation_func(class_data["activation_function"])
+        self.fitting_precision = get_precision(class_data['precision'])
+        self.useBN = False
+
+
+    def get_sel_type(self):
+        return self.sel_type
+
+    def get_wfc_numb(self):
+        return self.wfc_numb
+
+    def get_out_size(self):
+        return self.wfc_numb * 3
+
+    def build (self, 
+               input_d,
+               rot_mat,
+               natoms,
+               reuse = None,
+               suffix = '') :
+        start_index = 0
+        inputs = tf.cast(tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision)
+        rot_mat = tf.reshape(rot_mat, [-1, 9 * natoms[0]])
+
+        count = 0
+        for type_i in range(self.ntypes):
+            # cut-out inputs
+            inputs_i = tf.slice (inputs,
+                                 [ 0, start_index*      self.dim_descrpt],
+                                 [-1, natoms[2+type_i]* self.dim_descrpt] )
+            inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt])
+            rot_mat_i = tf.slice (rot_mat,
+                                  [ 0, start_index*      9],
+                                  [-1, natoms[2+type_i]* 9] )
+            rot_mat_i = tf.reshape(rot_mat_i, [-1, 3, 3])
+            start_index += natoms[2+type_i]
+            if not type_i in self.sel_type :
+                continue
+            layer = inputs_i
+            for ii in range(0,len(self.n_neuron)) :
+                if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] :
+                    layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed)
+                else :
+                    layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed)
+                if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+            # (nframes x natoms) x (nwfc x 3)
+            final_layer = one_layer(layer, self.wfc_numb * 3, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, uniform_seed = self.uniform_seed)
+            if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift
+            # (nframes x natoms) x nwfc(wc) x 3(coord_local)
+            final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.wfc_numb, 3])
+            # (nframes x natoms) x nwfc(wc) x 3(coord)
+            final_layer = tf.matmul(final_layer, rot_mat_i)
+            # nframes x natoms x nwfc(wc) x 3(coord_local)
+            final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], self.wfc_numb, 3])
+
+            # concat the results
+            if count == 0:
+                outs = final_layer
+            else:
+                outs = tf.concat([outs, final_layer], axis = 1)
+            count += 1
+        
+        tf.summary.histogram('fitting_net_output', outs)
+        return tf.cast(tf.reshape(outs, [-1]),  GLOBAL_TF_FLOAT_PRECISION)

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/infer/__init__.py

+"""Submodule containing all the implemented potentials."""
+
+from pathlib import Path
+from typing import Union
+
+from .data_modifier import DipoleChargeModifier
+from .deep_dipole import DeepDipole
+from .deep_eval import DeepEval
+from .deep_polar import DeepGlobalPolar, DeepPolar
+from .deep_pot import DeepPot
+from .deep_wfc import DeepWFC
+from .ewald_recp import EwaldRecp
+from .model_devi import calc_model_devi
+
+__all__ = [
+    "DeepPotential",
+    "DeepDipole",
+    "DeepEval",
+    "DeepGlobalPolar",
+    "DeepPolar",
+    "DeepPot",
+    "DeepWFC",
+    "DipoleChargeModifier",
+    "EwaldRecp",
+    "calc_model_devi"
+]
+
+
+def DeepPotential(
+    model_file: Union[str, Path],
+    load_prefix: str = "load",
+    default_tf_graph: bool = False,
+) -> Union[DeepDipole, DeepGlobalPolar, DeepPolar, DeepPot, DeepWFC]:
+    """Factory function that will inialize appropriate potential read from `model_file`.
+
+    Parameters
+    ----------
+    model_file: str
+        The name of the frozen model file.
+    load_prefix: str
+        The prefix in the load computational graph
+    default_tf_graph : bool
+        If uses the default tf graph, otherwise build a new tf graph for evaluation
+
+    Returns
+    -------
+    Union[DeepDipole, DeepGlobalPolar, DeepPolar, DeepPot, DeepWFC]
+        one of the available potentials
+
+    Raises
+    ------
+    RuntimeError
+        if model file does not correspond to any implementd potential
+    """
+    mf = Path(model_file)
+
+    model_type = DeepEval(
+        mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph
+    ).model_type
+
+    if model_type == "ener":
+        dp = DeepPot(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph)
+    elif model_type == "dipole":
+        dp = DeepDipole(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph)
+    elif model_type == "polar":
+        dp = DeepPolar(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph)
+    elif model_type == "global_polar":
+        dp = DeepGlobalPolar(
+            mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph
+        )
+    elif model_type == "wfc":
+        dp = DeepWFC(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph)
+    else:
+        raise RuntimeError(f"unknow model type {model_type}")
+
+    return dp

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/infer/data_modifier.py

+import os
+import numpy as np
+from typing import Tuple, List
+
+from deepmd.infer.deep_dipole import DeepDipole
+from deepmd.infer.ewald_recp import EwaldRecp
+from deepmd.env import tf
+from deepmd.common import select_idx_map, make_default_mesh
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
+from deepmd.env import GLOBAL_ENER_FLOAT_PRECISION
+from deepmd.env import global_cvt_2_tf_float
+from deepmd.env import global_cvt_2_ener_float
+from deepmd.env import op_module
+from deepmd.utils.sess import run_sess
+
+
+class DipoleChargeModifier(DeepDipole):
+    """
+    
+    Parameters
+    ----------
+    model_name
+            The model file for the DeepDipole model
+    model_charge_map
+            Gives the amount of charge for the wfcc
+    sys_charge_map
+            Gives the amount of charge for the real atoms
+    ewald_h
+            Grid spacing of the reciprocal part of Ewald sum. Unit: A
+    ewald_beta
+            Splitting parameter of the Ewald sum. Unit: A^{-1}
+    """
+    def __init__(self, 
+                 model_name : str, 
+                 model_charge_map : List[float],
+                 sys_charge_map : List[float], 
+                 ewald_h : float = 1, 
+                 ewald_beta : float = 1
+    ) -> None:
+        """
+        Constructor 
+        """
+        # the dipole model is loaded with prefix 'dipole_charge'
+        self.modifier_prefix = 'dipole_charge'
+        # init dipole model
+        DeepDipole.__init__(self, 
+                            model_name, 
+                            load_prefix = self.modifier_prefix, 
+                            default_tf_graph = True)
+        self.model_name = model_name
+        self.model_charge_map = model_charge_map
+        self.sys_charge_map = sys_charge_map
+        self.sel_type = list(self.get_sel_type())
+        # init ewald recp
+        self.ewald_h = ewald_h
+        self.ewald_beta = ewald_beta
+        self.er = EwaldRecp(self.ewald_h, self.ewald_beta)
+        # dimension of dipole
+        self.ext_dim = 3
+        self.t_ndesc  = self.graph.get_tensor_by_name(os.path.join(self.modifier_prefix, 'descrpt_attr/ndescrpt:0'))
+        self.t_sela  = self.graph.get_tensor_by_name(os.path.join(self.modifier_prefix, 'descrpt_attr/sel:0'))
+        [self.ndescrpt, self.sel_a] = run_sess(self.sess, [self.t_ndesc, self.t_sela])
+        self.sel_r = [ 0 for ii in range(len(self.sel_a)) ]
+        self.nnei_a = np.cumsum(self.sel_a)[-1]
+        self.nnei_r = np.cumsum(self.sel_r)[-1]
+        self.nnei = self.nnei_a + self.nnei_r
+        self.ndescrpt_a = self.nnei_a * 4
+        self.ndescrpt_r = self.nnei_r * 1
+        assert(self.ndescrpt == self.ndescrpt_a + self.ndescrpt_r)
+        self.force = None
+        self.ntypes = len(self.sel_a)
+
+
+    def build_fv_graph(self) -> tf.Tensor:
+        """
+        Build the computational graph for the force and virial inference.
+        """
+        with tf.variable_scope('modifier_attr') :
+            t_mdl_name = tf.constant(self.model_name, 
+                                     name = 'mdl_name', 
+                                     dtype = tf.string)
+            t_modi_type = tf.constant(self.modifier_prefix, 
+                                      name = 'type', 
+                                      dtype = tf.string)
+            t_mdl_charge_map = tf.constant(' '.join([str(ii) for ii in self.model_charge_map]),
+                                            name = 'mdl_charge_map', 
+                                            dtype = tf.string)
+            t_sys_charge_map = tf.constant(' '.join([str(ii) for ii in self.sys_charge_map]),
+                                            name = 'sys_charge_map', 
+                                            dtype = tf.string)
+            t_ewald_h = tf.constant(self.ewald_h,
+                                    name = 'ewald_h', 
+                                    dtype = tf.float64)
+            t_ewald_b = tf.constant(self.ewald_beta,
+                                    name = 'ewald_beta',
+                                    dtype = tf.float64)
+        with self.graph.as_default():
+            return self._build_fv_graph_inner()        
+
+
+    def _build_fv_graph_inner(self):
+        self.t_ef = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name = 't_ef')
+        nf = 10
+        nfxnas = 64*nf
+        nfxna = 192*nf
+        nf = -1
+        nfxnas = -1
+        nfxna = -1
+        self.t_box_reshape = tf.reshape(self.t_box, [-1, 9])
+        t_nframes = tf.shape(self.t_box_reshape)[0]
+
+        # (nframes x natoms) x ndescrpt
+        self.descrpt = self.graph.get_tensor_by_name(os.path.join(self.modifier_prefix, 'o_rmat:0'))
+        self.descrpt_deriv = self.graph.get_tensor_by_name(os.path.join(self.modifier_prefix, 'o_rmat_deriv:0'))
+        self.nlist = self.graph.get_tensor_by_name(os.path.join(self.modifier_prefix, 'o_nlist:0'))
+        self.rij = self.graph.get_tensor_by_name(os.path.join(self.modifier_prefix, 'o_rij:0'))
+        # self.descrpt_reshape = tf.reshape(self.descrpt, [nf, 192 * self.ndescrpt])
+        # self.descrpt_deriv = tf.reshape(self.descrpt_deriv, [nf, 192 * self.ndescrpt * 3])
+
+        # nframes x (natoms_sel x 3)
+        self.t_ef_reshape = tf.reshape(self.t_ef,  [t_nframes, -1])
+        # nframes x (natoms x 3)
+        self.t_ef_reshape = self._enrich(self.t_ef_reshape, dof = 3)
+        # (nframes x natoms) x 3
+        self.t_ef_reshape = tf.reshape(self.t_ef_reshape, [nfxna, 3])
+        # nframes x (natoms_sel x 3)
+        self.t_tensor_reshape = tf.reshape(self.t_tensor, [t_nframes, -1])
+        # nframes x (natoms x 3)
+        self.t_tensor_reshape = self._enrich(self.t_tensor_reshape, dof = 3)
+        # (nframes x natoms) x 3
+        self.t_tensor_reshape = tf.reshape(self.t_tensor_reshape, [nfxna, 3])
+        # (nframes x natoms) x ndescrpt
+        [self.t_ef_d] = tf.gradients(self.t_tensor_reshape, self.descrpt, self.t_ef_reshape)
+        # nframes x (natoms x ndescrpt)
+        self.t_ef_d = tf.reshape(self.t_ef_d, [nf, self.t_natoms[0] * self.ndescrpt])
+        # t_ef_d is force (with -1), prod_forc takes deriv, so we need the opposite
+        self.t_ef_d_oppo = -self.t_ef_d
+        
+        force = op_module.prod_force_se_a(self.t_ef_d_oppo,
+                                          self.descrpt_deriv, 
+                                          self.nlist, 
+                                          self.t_natoms,
+                                          n_a_sel = self.nnei_a,
+                                          n_r_sel = self.nnei_r)
+        virial, atom_virial \
+            = op_module.prod_virial_se_a (self.t_ef_d_oppo,
+                                          self.descrpt_deriv,
+                                          self.rij,
+                                          self.nlist,
+                                          self.t_natoms,
+                                          n_a_sel = self.nnei_a,
+                                          n_r_sel = self.nnei_r)
+        force = tf.identity(force, name='o_dm_force')
+        virial = tf.identity(virial, name='o_dm_virial')
+        atom_virial = tf.identity(atom_virial, name='o_dm_av')
+        return force, virial, atom_virial
+
+
+    def _enrich(self, dipole, dof = 3):
+        coll = []                
+        sel_start_idx = 0
+        for type_i in range(self.ntypes):
+            if type_i in self.sel_type:
+                di = tf.slice(dipole, 
+                              [ 0, sel_start_idx           * dof],
+                              [-1, self.t_natoms[2+type_i] * dof])
+                sel_start_idx += self.t_natoms[2+type_i]
+            else:
+                di = tf.zeros([tf.shape(dipole)[0], self.t_natoms[2+type_i] * dof],
+                              dtype = GLOBAL_TF_FLOAT_PRECISION)
+            coll.append(di)
+        return tf.concat(coll, axis = 1)
+
+    def _slice_descrpt_deriv(self, deriv):
+        coll = []
+        start_idx = 0
+        for type_i in range(self.ntypes):
+            if type_i in self.sel_type:
+                di = tf.slice(deriv, 
+                              [ 0, start_idx               * self.ndescrpt],
+                              [-1, self.t_natoms[2+type_i] * self.ndescrpt])
+                coll.append(di)
+            start_idx += self.t_natoms[2+type_i]
+        return tf.concat(coll, axis = 1)        
+
+
+    def eval(self, 
+             coord : np.ndarray, 
+             box : np.ndarray, 
+             atype : np.ndarray, 
+             eval_fv : bool = True
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Evaluate the modification
+        
+        Parameters
+        ----------
+        coord
+                The coordinates of atoms
+        box
+                The simulation region. PBC is assumed
+        atype
+                The atom types
+        eval_fv
+                Evaluate force and virial
+
+        Returns
+        -------
+        tot_e
+                The energy modification
+        tot_f
+                The force modification
+        tot_v
+                The virial modification
+        """
+        atype = np.array(atype, dtype=int)
+        coord, atype, imap = self.sort_input(coord, atype)
+        # natoms = coord.shape[1] // 3
+        natoms = atype.size
+        nframes = coord.shape[0]
+        box = np.reshape(box, [nframes, 9])
+        atype = np.reshape(atype, [natoms])
+        sel_idx_map = select_idx_map(atype, self.sel_type)
+        nsel = len(sel_idx_map)
+        # setup charge
+        charge = np.zeros([natoms])
+        for ii in range(natoms):
+            charge[ii] = self.sys_charge_map[atype[ii]]
+        charge = np.tile(charge, [nframes, 1])
+
+        # add wfcc
+        all_coord, all_charge, dipole = self._extend_system(coord, box, atype, charge)
+        
+        # print('compute er')
+        batch_size = 5
+        tot_e = []
+        all_f = []
+        all_v = []
+        for ii in range(0,nframes,batch_size):
+            e,f,v = self.er.eval(all_coord[ii:ii+batch_size], all_charge[ii:ii+batch_size], box[ii:ii+batch_size])
+            tot_e.append(e)
+            all_f.append(f)
+            all_v.append(v)
+        tot_e = np.concatenate(tot_e, axis = 0)
+        all_f = np.concatenate(all_f, axis = 0)
+        all_v = np.concatenate(all_v, axis = 0)
+        # print('finish  er')
+        # reshape
+        tot_e.reshape([nframes,1])
+
+        tot_f = None
+        tot_v = None
+        if self.force is None:
+            self.force, self.virial, self.av = self.build_fv_graph()
+        if eval_fv:
+            # compute f
+            ext_f = all_f[:,natoms*3:]
+            corr_f = []
+            corr_v = []
+            corr_av = []
+            for ii in range(0,nframes,batch_size):
+                f, v, av = self._eval_fv(coord[ii:ii+batch_size], box[ii:ii+batch_size], atype, ext_f[ii:ii+batch_size])
+                corr_f.append(f)
+                corr_v.append(v)
+                corr_av.append(av)
+            corr_f = np.concatenate(corr_f, axis = 0)
+            corr_v = np.concatenate(corr_v, axis = 0)
+            corr_av = np.concatenate(corr_av, axis = 0)
+            tot_f = all_f[:,:natoms*3] + corr_f
+            for ii in range(nsel):            
+                orig_idx = sel_idx_map[ii]            
+                tot_f[:,orig_idx*3:orig_idx*3+3] += ext_f[:,ii*3:ii*3+3]                
+            tot_f = self.reverse_map(np.reshape(tot_f, [nframes,-1,3]), imap)
+            # reshape
+            tot_f = tot_f.reshape([nframes,natoms,3])
+            # compute v
+            dipole3 = np.reshape(dipole, [nframes, nsel, 3])
+            ext_f3 = np.reshape(ext_f, [nframes, nsel, 3])
+            ext_f3 = np.transpose(ext_f3, [0, 2, 1])
+            # fd_corr_v = -np.matmul(ext_f3, dipole3).T.reshape([nframes, 9])
+            # fd_corr_v = -np.matmul(ext_f3, dipole3)
+            # fd_corr_v = np.transpose(fd_corr_v, [0, 2, 1]).reshape([nframes, 9])
+            fd_corr_v = -np.matmul(ext_f3, dipole3).reshape([nframes, 9])
+            # print(all_v, '\n', corr_v, '\n', fd_corr_v)
+            tot_v = all_v + corr_v + fd_corr_v
+            # reshape
+            tot_v = tot_v.reshape([nframes,9])
+
+        return tot_e, tot_f, tot_v
+
+
+    def _eval_fv(self, coords, cells, atom_types, ext_f) :
+        # reshape the inputs 
+        cells = np.reshape(cells, [-1, 9])
+        nframes = cells.shape[0]
+        coords = np.reshape(coords, [nframes, -1])
+        natoms = coords.shape[1] // 3
+
+        # sort inputs
+        coords, atom_types, imap, sel_at, sel_imap = self.sort_input(coords, atom_types, sel_atoms = self.get_sel_type())
+
+        # make natoms_vec and default_mesh
+        natoms_vec = self.make_natoms_vec(atom_types)
+        assert(natoms_vec[0] == natoms)
+        default_mesh = make_default_mesh(cells)
+
+        # evaluate
+        tensor = []
+        feed_dict_test = {}
+        feed_dict_test[self.t_natoms] = natoms_vec
+        feed_dict_test[self.t_type  ] = np.tile(atom_types, [nframes, 1]).reshape([-1])
+        feed_dict_test[self.t_coord ] = coords.reshape([-1])
+        feed_dict_test[self.t_box   ] = cells.reshape([-1])
+        feed_dict_test[self.t_mesh  ] = default_mesh.reshape([-1])
+        feed_dict_test[self.t_ef    ] = ext_f.reshape([-1])
+        # print(run_sess(self.sess, tf.shape(self.t_tensor), feed_dict = feed_dict_test))
+        fout, vout, avout \
+            = run_sess(self.sess, [self.force, self.virial, self.av],
+                            feed_dict = feed_dict_test)
+        # print('fout: ', fout.shape, fout)
+        fout = self.reverse_map(np.reshape(fout, [nframes,-1,3]), imap)
+        fout = np.reshape(fout, [nframes, -1])
+        return fout, vout, avout
+
+
+    def _extend_system(self, coord, box, atype, charge):
+        natoms = coord.shape[1] // 3
+        nframes = coord.shape[0]
+        # sel atoms and setup ref coord
+        sel_idx_map = select_idx_map(atype, self.sel_type)
+        nsel = len(sel_idx_map)
+        coord3 = coord.reshape([nframes, natoms, 3])
+        ref_coord = coord3[:,sel_idx_map,:]
+        ref_coord = np.reshape(ref_coord, [nframes, nsel * 3])
+        
+        batch_size = 8
+        all_dipole = []
+        for ii in range(0,nframes,batch_size):
+            dipole = DeepDipole.eval(self,
+                                     coord[ii:ii+batch_size],
+                                     box[ii:ii+batch_size],
+                                     atype)
+            all_dipole.append(dipole)
+        dipole = np.concatenate(all_dipole, axis = 0)
+        assert(dipole.shape[0] == nframes)
+        dipole = np.reshape(dipole, [nframes, nsel * 3])
+        
+        wfcc_coord = ref_coord + dipole
+        # wfcc_coord = dipole
+        wfcc_charge = np.zeros([nsel])
+        for ii in range(nsel):
+            orig_idx = self.sel_type.index(atype[sel_idx_map[ii]])
+            wfcc_charge[ii] = self.model_charge_map[orig_idx]
+        wfcc_charge = np.tile(wfcc_charge, [nframes, 1])
+
+        wfcc_coord = np.reshape(wfcc_coord, [nframes, nsel * 3])
+        wfcc_charge = np.reshape(wfcc_charge, [nframes, nsel])
+
+        all_coord = np.concatenate((coord, wfcc_coord), axis = 1)
+        all_charge = np.concatenate((charge, wfcc_charge), axis = 1)
+
+        return all_coord, all_charge, dipole
+
+
+    def modify_data(self, 
+                    data : dict) -> None:        
+        """
+        Modify data.
+
+        Parameters
+        ----------
+        data
+                Internal data of DeepmdData.
+                Be a dict, has the following keys
+                - coord         coordinates
+                - box           simulation box
+                - type          atom types
+                - find_energy   tells if data has energy
+                - find_force    tells if data has force
+                - find_virial   tells if data has virial
+                - energy        energy
+                - force         force
+                - virial        virial
+        """
+        if 'find_energy' not in data and 'find_force' not in data and 'find_virial' not in data:
+            return
+
+        get_nframes=None
+        coord = data['coord'][:get_nframes,:]
+        box = data['box'][:get_nframes,:]
+        atype = data['type'][:get_nframes,:]
+        atype = atype[0]
+        nframes = coord.shape[0]
+
+        tot_e, tot_f, tot_v = self.eval(coord, box, atype)
+
+        # print(tot_f[:,0])
+        
+        if 'find_energy' in data and data['find_energy'] == 1.0 :
+            data['energy'] -= tot_e.reshape(data['energy'].shape)
+        if 'find_force' in data and data['find_force'] == 1.0 :
+            data['force'] -= tot_f.reshape(data['force'].shape)
+        if 'find_virial' in data and data['find_virial'] == 1.0 :
+            data['virial'] -= tot_v.reshape(data['virial'].shape)
+
+
+                           

_skbuild/linux-x86_64-3.6/setuptools/lib/deepmd/infer/deep_dipole.py

+from deepmd.infer.deep_tensor import DeepTensor
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+class DeepDipole(DeepTensor):
+    """Constructor.
+
+    Parameters
+    ----------
+    model_file : Path
+        The name of the frozen model file.
+    load_prefix: str
+        The prefix in the load computational graph
+    default_tf_graph : bool
+        If uses the default tf graph, otherwise build a new tf graph for evaluation
+
+    Warnings
+    --------
+    For developers: `DeepTensor` initializer must be called at the end after
+    `self.tensors` are modified because it uses the data in `self.tensors` dict.
+    Do not chanage the order!
+    """
+
+    def __init__(
+        self, model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False
+    ) -> None:
+
+        # use this in favor of dict update to move attribute from class to
+        # instance namespace
+        self.tensors = dict(
+            {
+                # output tensor
+                "t_tensor": "o_dipole:0",
+            },
+            **self.tensors
+        )
+
+        DeepTensor.__init__(
+            self,
+            model_file,
+            load_prefix=load_prefix,
+            default_tf_graph=default_tf_graph,
+        )
+
+    def get_dim_fparam(self) -> int:
+        """Unsupported in this model."""
+        raise NotImplementedError("This model type does not support this attribute")
+
+    def get_dim_aparam(self) -> int:
+        """Unsupported in this model."""
+        raise NotImplementedError("This model type does not support this attribute")