From 7a6fc98d98f73af20c7385ca914b351439abb46f Mon Sep 17 00:00:00 2001 From: birfied Date: Sun, 14 Sep 2025 12:41:02 +0800 Subject: [PATCH 1/2] docs: addmissing docs for mindchemistry --- .../crystalflow/models/conditioning.py | 2 +- .../applications/crystalflow/models/cspnet.py | 6 ++-- .../crystalflow/models/cspnet_condition.py | 6 ++-- .../crystalflow/models/diff_utils.py | 2 +- .../applications/crystalflow/models/flow.py | 4 +-- .../crystalflow/models/flow_condition.py | 4 +-- .../crystalflow/models/infer_utils.py | 2 +- .../applications/diffcsp/models/cspnet.py | 6 ++-- .../applications/diffcsp/models/diff_utils.py | 2 +- .../applications/diffcsp/models/diffusion.py | 4 +-- .../diffcsp/models/infer_utils.py | 2 +- MindChemistry/applications/orb/src/base.py | 8 ++--- .../orb/src/featurization_utilities.py | 8 ++--- .../mindchemistry/cell/allegro/allegro.py | 2 +- MindChemistry/mindchemistry/cell/cspnet.py | 6 ++-- .../mindchemistry/cell/dimenet/preprocess.py | 2 +- .../mindchemistry/cell/gemnet/gemnet_wrap.py | 2 +- MindChemistry/mindchemistry/cell/orb/gns.py | 4 +-- MindChemistry/mindchemistry/cell/orb/orb.py | 8 ++--- MindChemistry/mindchemistry/cell/orb/utils.py | 2 +- .../cell/mindchemistry.cell.Allegro.rst | 14 +++++++- .../cell/mindchemistry.cell.AutoEncoder.rst | 33 +++++++++++++++++++ .../cell/mindchemistry.cell.FCNet.rst | 10 ++++++ .../mindchemistry.cell.orb.EnergyHead.rst | 17 +++++++++- .../cell/mindchemistry.cell.orb.GraphHead.rst | 18 +++++++++- .../mindchemistry.cell.orb.MoleculeGNS.rst | 22 ++++++++++++- .../cell/mindchemistry.cell.orb.NodeHead.rst | 13 +++++++- .../cell/mindchemistry.cell.orb.Orb.rst | 21 +++++++++++- 28 files changed, 183 insertions(+), 47 deletions(-) diff --git a/MindChemistry/applications/crystalflow/models/conditioning.py b/MindChemistry/applications/crystalflow/models/conditioning.py index 799b8b6ad..da21f09c5 100644 --- a/MindChemistry/applications/crystalflow/models/conditioning.py +++ b/MindChemistry/applications/crystalflow/models/conditioning.py @@ -26,7 +26,7 @@ class GaussianExpansion(nn.Cell): Args: prop (Tensor): property values of (N_b x 1) shape. Returns: - Tensor: layer output of (N_b x N_g) shape. + Tensor. layer output of (N_b x N_g) shape. """ prop = prop.reshape(prop.shape[0], -1) coeff = -0.5 / ops.pow(self.widths, 2)[None, :] diff --git a/MindChemistry/applications/crystalflow/models/cspnet.py b/MindChemistry/applications/crystalflow/models/cspnet.py index c5ae68956..2efd7763b 100644 --- a/MindChemistry/applications/crystalflow/models/cspnet.py +++ b/MindChemistry/applications/crystalflow/models/cspnet.py @@ -51,7 +51,7 @@ class SinusoidsEmbedding(nn.Cell): x (Tensor): Distance Returns: - Tensor: Fourier embedding + Tensor. Fourier embedding """ emb = ops.ExpandDims()(x, -1) * Tensor( self.frequencies, @@ -70,7 +70,7 @@ def mul_mask(features, mask): and value 0 specifies the corresponding dimension of input tensor to be zero. Returns: - Tensor: Output tensor + Tensor. Output tensor """ return ops.mul(features, ops.reshape(mask, (-1, 1))) @@ -172,7 +172,7 @@ class CSPLayer(nn.Cell): edge_mask (Tensor): Edge mask for padded tensor. Shape: (num_edges,) Returns: - Tensor: The output tensor. Shape: (num_atoms, hidden_dim) + Tensor. The output tensor. Shape: (num_atoms, hidden_dim) """ node_input = node_features node_features = mul_mask(node_features, node_mask) diff --git a/MindChemistry/applications/crystalflow/models/cspnet_condition.py b/MindChemistry/applications/crystalflow/models/cspnet_condition.py index 9ae1c2cdd..791ea4140 100644 --- a/MindChemistry/applications/crystalflow/models/cspnet_condition.py +++ b/MindChemistry/applications/crystalflow/models/cspnet_condition.py @@ -51,7 +51,7 @@ class SinusoidsEmbedding(nn.Cell): x (Tensor): Distance Returns: - Tensor: Fourier embedding + Tensor. Fourier embedding """ emb = ops.ExpandDims()(x, -1) * Tensor( self.frequencies, @@ -70,7 +70,7 @@ def mul_mask(features, mask): and value 0 specifies the corresponding dimension of input tensor to be zero. Returns: - Tensor: Output tensor + Tensor. Output tensor """ return ops.mul(features, ops.reshape(mask, (-1, 1))) @@ -173,7 +173,7 @@ class CSPLayer(nn.Cell): edge_mask (Tensor): Edge mask for padded tensor. Shape: (num_edges,) Returns: - Tensor: The output tensor. Shape: (num_atoms, hidden_dim) + Tensor. The output tensor. Shape: (num_atoms, hidden_dim) """ node_input = node_features node_features = mul_mask(node_features, node_mask) diff --git a/MindChemistry/applications/crystalflow/models/diff_utils.py b/MindChemistry/applications/crystalflow/models/diff_utils.py index e74e20a00..7eff1cdf0 100644 --- a/MindChemistry/applications/crystalflow/models/diff_utils.py +++ b/MindChemistry/applications/crystalflow/models/diff_utils.py @@ -112,7 +112,7 @@ def d_log_p_wrapped_normal_ms(x, sigma, n=10, t=1.0): t (int): The period of wrapped normal distribution. Defaults to 1.0. Returns: - Tensor: The score for the input value x. + Tensor. The score for the input value x. """ p_ = 0 for i in range(-n, n + 1): diff --git a/MindChemistry/applications/crystalflow/models/flow.py b/MindChemistry/applications/crystalflow/models/flow.py index 821f3520b..6098c5f8d 100644 --- a/MindChemistry/applications/crystalflow/models/flow.py +++ b/MindChemistry/applications/crystalflow/models/flow.py @@ -47,7 +47,7 @@ class SinusoidalTimeEmbeddings(nn.Cell): time (Tensor): flow time step Returns: - Tensor: Time embedding + Tensor. Time embedding """ half_dim = self.dim // 2 embeddings = math.log(10000) / (half_dim - 1) @@ -65,7 +65,7 @@ def lattice_params_to_matrix_mindspore(lengths, angles): lengths (Tensor): Tensor of shape (N, 3), unit A angles (Tensor):: Tensor of shape (N, 3), unit degree Returns: - Tensor: Tensor of shape (N, 3, 3) + Tensor. Tensor of shape (N, 3, 3) """ angles_r = ops.deg2rad(angles) coses = ops.cos(angles_r) diff --git a/MindChemistry/applications/crystalflow/models/flow_condition.py b/MindChemistry/applications/crystalflow/models/flow_condition.py index 938012b7c..f3731f80d 100644 --- a/MindChemistry/applications/crystalflow/models/flow_condition.py +++ b/MindChemistry/applications/crystalflow/models/flow_condition.py @@ -48,7 +48,7 @@ class SinusoidalTimeEmbeddings(nn.Cell): time (Tensor): flow time step Returns: - Tensor: Time embedding + Tensor. Time embedding """ half_dim = self.dim // 2 embeddings = math.log(10000) / (half_dim - 1) @@ -66,7 +66,7 @@ def lattice_params_to_matrix_mindspore(lengths, angles): lengths (Tensor): Tensor of shape (N, 3), unit A angles (Tensor):: Tensor of shape (N, 3), unit degree Returns: - Tensor: Tensor of shape (N, 3, 3) + Tensor. Tensor of shape (N, 3, 3) """ angles_r = ops.deg2rad(angles) coses = ops.cos(angles_r) diff --git a/MindChemistry/applications/crystalflow/models/infer_utils.py b/MindChemistry/applications/crystalflow/models/infer_utils.py index 7ff160955..2c72b3e07 100644 --- a/MindChemistry/applications/crystalflow/models/infer_utils.py +++ b/MindChemistry/applications/crystalflow/models/infer_utils.py @@ -50,7 +50,7 @@ def count_consecutive_occurrences(lst): lst (list): The input list Returns: - list: List of numbers of consecutive occurrences of each digit in the list. + list. List of numbers of consecutive occurrences of each digit in the list. """ if not lst: return [] diff --git a/MindChemistry/applications/diffcsp/models/cspnet.py b/MindChemistry/applications/diffcsp/models/cspnet.py index a12731b41..363c8d1c3 100644 --- a/MindChemistry/applications/diffcsp/models/cspnet.py +++ b/MindChemistry/applications/diffcsp/models/cspnet.py @@ -51,7 +51,7 @@ class SinusoidsEmbedding(nn.Cell): x (Tensor): Distance Returns: - Tensor: Fourier embedding + Tensor. Fourier embedding """ emb = ops.ExpandDims()(x, -1) * Tensor( self.frequencies, @@ -70,7 +70,7 @@ def mul_mask(features, mask): and value 0 specifies the corresponding dimension of input tensor to be zero. Returns: - Tensor: Output tensor + Tensor. Output tensor """ return ops.mul(features, ops.reshape(mask, (-1, 1))) @@ -175,7 +175,7 @@ class CSPLayer(nn.Cell): edge_mask (Tensor): Edge mask for padded tensor. Shape: (num_edges,) Returns: - Tensor: The output tensor. Shape: (num_atoms, hidden_dim) + Tensor. The output tensor. Shape: (num_atoms, hidden_dim) """ node_input = node_features node_features = mul_mask(node_features, node_mask) diff --git a/MindChemistry/applications/diffcsp/models/diff_utils.py b/MindChemistry/applications/diffcsp/models/diff_utils.py index e74e20a00..7eff1cdf0 100644 --- a/MindChemistry/applications/diffcsp/models/diff_utils.py +++ b/MindChemistry/applications/diffcsp/models/diff_utils.py @@ -112,7 +112,7 @@ def d_log_p_wrapped_normal_ms(x, sigma, n=10, t=1.0): t (int): The period of wrapped normal distribution. Defaults to 1.0. Returns: - Tensor: The score for the input value x. + Tensor. The score for the input value x. """ p_ = 0 for i in range(-n, n + 1): diff --git a/MindChemistry/applications/diffcsp/models/diffusion.py b/MindChemistry/applications/diffcsp/models/diffusion.py index 21577030c..7b29a861c 100644 --- a/MindChemistry/applications/diffcsp/models/diffusion.py +++ b/MindChemistry/applications/diffcsp/models/diffusion.py @@ -48,7 +48,7 @@ class SinusoidalTimeEmbeddings(nn.Cell): time (Tensor): diffution time step Returns: - Tensor: Time embedding + Tensor. Time embedding """ half_dim = self.dim // 2 embeddings = math.log(10000) / (half_dim - 1) @@ -66,7 +66,7 @@ def lattice_params_to_matrix_mindspore(lengths, angles): lengths (Tensor): Tensor of shape (N, 3), unit A angles (Tensor):: Tensor of shape (N, 3), unit degree Returns: - Tensor: Tensor of shape (N, 3, 3) + Tensor. Tensor of shape (N, 3, 3) """ angles_r = ops.deg2rad(angles) coses = ops.cos(angles_r) diff --git a/MindChemistry/applications/diffcsp/models/infer_utils.py b/MindChemistry/applications/diffcsp/models/infer_utils.py index 7ff160955..2c72b3e07 100644 --- a/MindChemistry/applications/diffcsp/models/infer_utils.py +++ b/MindChemistry/applications/diffcsp/models/infer_utils.py @@ -50,7 +50,7 @@ def count_consecutive_occurrences(lst): lst (list): The input list Returns: - list: List of numbers of consecutive occurrences of each digit in the list. + list. List of numbers of consecutive occurrences of each digit in the list. """ if not lst: return [] diff --git a/MindChemistry/applications/orb/src/base.py b/MindChemistry/applications/orb/src/base.py index 046e5950f..861fd38d0 100644 --- a/MindChemistry/applications/orb/src/base.py +++ b/MindChemistry/applications/orb/src/base.py @@ -172,7 +172,7 @@ class AtomGraphs(NamedTuple): Any None attributes are not included in the dictionary. Returns: - dict: A dictionary mapping attribute_name -> tensor/scalar + dict. A dictionary mapping attribute_name -> tensor/scalar """ ret = {} for key, val in self._asdict().items(): @@ -190,7 +190,7 @@ class AtomGraphs(NamedTuple): """Return a single dictionary mapping each AtomGraph property to a corresponding list of tensors/scalars. Returns: - dict: A dict mapping attribute_name -> list of length batch_size containing tensors/scalars. + dict. A dict mapping attribute_name -> list of length batch_size containing tensors/scalars. """ batch_dict = defaultdict(list) for graph in self.split(self): @@ -272,7 +272,7 @@ def batch_graphs(graphs: List[AtomGraphs]) -> AtomGraphs: graphs (List[AtomGraphs]): A list of AtomGraphs to be batched together. Returns: - AtomGraphs: A new AtomGraphs object with the concatenated nodes, + AtomGraphs. A new AtomGraphs object with the concatenated nodes, edges, and features from the input graphs, along with concatenated target, system ID, and other information. """ @@ -342,7 +342,7 @@ def refeaturize_atomgraphs( which considers boundary crossing for more than cell is used. Returns: - AtomGraphs: A refeaturized AtomGraphs object. + AtomGraphs. A refeaturized AtomGraphs object. """ if cell is None: cell = atoms.cell diff --git a/MindChemistry/applications/orb/src/featurization_utilities.py b/MindChemistry/applications/orb/src/featurization_utilities.py index dd68d4ad8..a85f3d06e 100644 --- a/MindChemistry/applications/orb/src/featurization_utilities.py +++ b/MindChemistry/applications/orb/src/featurization_utilities.py @@ -42,7 +42,7 @@ def gaussian_basis_function( scale (Tensor, optional): The width of the gaussians. Defaults to 1. Returns: - Tensor: A tensor of shape [num_scalars, num_bases]. + Tensor. A tensor of shape [num_scalars, num_bases]. """ assert len(scalars.shape) == 1 gaussian_means = ops.arange( @@ -63,7 +63,7 @@ def featurize_edges( distance_featurization (DistanceFeaturization): A function that featurizes the distances of the vectors. Returns: - tensor: Edge features, shape [num_edge, num_edge_features]. + tensor. Edge features, shape [num_edge, num_edge_features]. """ edge_features = [] edge_norms = mint.linalg.norm(edge_vectors, dim=1) @@ -86,7 +86,7 @@ def compute_edge_vectors( positions (tensor): Positions of each node. Shape [num_nodes, 3] Returns: - tensor: The vectors of each edge. + tensor. The vectors of each edge. """ senders = edge_index[0] receivers = edge_index[1] @@ -160,7 +160,7 @@ def _compute_img_positions_torch( for each atom, i.e this function is agnostic as to whether it is computing with respect to a batch or not. Returns: - Tensor: The positions of the periodic images. Shape [num_atoms, 27, 3]. + Tensor. The positions of the periodic images. Shape [num_atoms, 27, 3]. """ num_positions = len(positions) diff --git a/MindChemistry/mindchemistry/cell/allegro/allegro.py b/MindChemistry/mindchemistry/cell/allegro/allegro.py index 741254e9d..2e6fbe646 100644 --- a/MindChemistry/mindchemistry/cell/allegro/allegro.py +++ b/MindChemistry/mindchemistry/cell/allegro/allegro.py @@ -34,7 +34,7 @@ def tp_path_exists(irreps_in1, irreps_in2, ir_out): ir_out: ir_out Returns: - bool: tp_path_exists or not + bool. tp_path_exists or not """ irreps_in1 = Irreps(irreps_in1).simplify() irreps_in2 = Irreps(irreps_in2).simplify() diff --git a/MindChemistry/mindchemistry/cell/cspnet.py b/MindChemistry/mindchemistry/cell/cspnet.py index be8a8a3e9..41db50948 100644 --- a/MindChemistry/mindchemistry/cell/cspnet.py +++ b/MindChemistry/mindchemistry/cell/cspnet.py @@ -57,7 +57,7 @@ class SinusoidsEmbedding(nn.Cell): x (Tensor): Distance Returns: - Tensor: Fourier embedding + Tensor. Fourier embedding """ emb = self.expand_dims(x, -1) * self.frequencies_tensor emb = self.reshape(emb, (-1, self.n_frequencies * self.n_space)) @@ -74,7 +74,7 @@ def mul_mask(features, mask): and value 0 specifies the corresponding dimension of input tensor to be zero. Returns: - Tensor: Output tensor + Tensor. Output tensor """ return ops.mul(features, ops.reshape(mask, (-1, 1))) @@ -180,7 +180,7 @@ class CSPLayer(nn.Cell): edge_mask (Tensor): Edge mask for padded tensor. Shape: (num_edges,) Returns: - Tensor: The output tensor. Shape: (num_atoms, hidden_dim) + Tensor. The output tensor. Shape: (num_atoms, hidden_dim) """ node_input = node_features node_features = mul_mask(node_features, node_mask) diff --git a/MindChemistry/mindchemistry/cell/dimenet/preprocess.py b/MindChemistry/mindchemistry/cell/dimenet/preprocess.py index 905ce82c1..48476e670 100644 --- a/MindChemistry/mindchemistry/cell/dimenet/preprocess.py +++ b/MindChemistry/mindchemistry/cell/dimenet/preprocess.py @@ -146,7 +146,7 @@ class PreProcess: edge_index (np.ndarray): Array of edge indices. Returns: - Tuple: Tuple containing computed triplets. + Tuple. Tuple containing computed triplets. """ node_1, node_2 = edge_index num_atoms = np.max(edge_index) + 1 diff --git a/MindChemistry/mindchemistry/cell/gemnet/gemnet_wrap.py b/MindChemistry/mindchemistry/cell/gemnet/gemnet_wrap.py index b769d8496..57857f302 100644 --- a/MindChemistry/mindchemistry/cell/gemnet/gemnet_wrap.py +++ b/MindChemistry/mindchemistry/cell/gemnet/gemnet_wrap.py @@ -107,7 +107,7 @@ class GemNetWrap: batch_size (int): The batch size. Returns: - tuple: A tuple containing the energy and force predictions. + tuple. A tuple containing the energy and force predictions. - E_a (ms.Tensor): The predicted energies. - F_j (ms.Tensor): The predicted forces. """ diff --git a/MindChemistry/mindchemistry/cell/orb/gns.py b/MindChemistry/mindchemistry/cell/orb/gns.py index ab53083f8..e489780b3 100644 --- a/MindChemistry/mindchemistry/cell/orb/gns.py +++ b/MindChemistry/mindchemistry/cell/orb/gns.py @@ -59,7 +59,7 @@ def get_cutoff(p: int, r: Tensor, r_max: float) -> Tensor: r_max (float): Maximum distance for the cutoff. Returns: - Tensor: Cutoff tensor. + Tensor. Cutoff tensor. """ envelope = 1.0 - ((p + 1.0) * (p + 2.0) / 2.0) * ops.pow(r / r_max, p) + \ p * (p + 2.0) * ops.pow(r / r_max, p + 1) - \ @@ -84,7 +84,7 @@ def gaussian_basis_function( scale (Tensor, optional): The width of the gaussians. Defaults to 1. Returns: - Tensor: A tensor of shape [num_scalars, num_bases]. + Tensor. A tensor of shape [num_scalars, num_bases]. """ assert len(scalars.shape) == 1 gaussian_means = ops.arange( diff --git a/MindChemistry/mindchemistry/cell/orb/orb.py b/MindChemistry/mindchemistry/cell/orb/orb.py index 8afc55dbf..404223e36 100644 --- a/MindChemistry/mindchemistry/cell/orb/orb.py +++ b/MindChemistry/mindchemistry/cell/orb/orb.py @@ -259,7 +259,7 @@ class NodeHead(ms.nn.Cell): n_node: Number of nodes in the graph. Returns: - node_pred: Node-level predictions of shape (n_nodes, target_property_dim). + node_pred. Node-level predictions of shape (n_nodes, target_property_dim). """ out = self(node_features, n_node) pred = out["node_pred"] @@ -378,7 +378,7 @@ class GraphHead(ms.nn.Cell): atomic_numbers: Optional atomic numbers for reference energy calculation Returns: - probs: Graph-level predictions of shape (n_graphs, target_property_dim). + probs. Graph-level predictions of shape (n_graphs, target_property_dim). If compute_stress is True, this will be the stress tensor. If compute_stress is False, this will be the graph-level property (e.g., energy). """ @@ -495,7 +495,7 @@ class EnergyHead(GraphHead): atomic_numbers: Optional atomic numbers for reference energy calculation Returns: - graph_pred: Energy prediction + graph_pred. Energy prediction """ if atomic_numbers is None: raise ValueError("atomic_numbers is required for energy prediction") @@ -671,7 +671,7 @@ class Orb(ms.nn.Cell): atomic_numbers: A tensor of atomic numbers for reference energy calculation. Returns: - ouput_dict: A dictionary containing the predictions: + ouput_dict. A dictionary containing the predictions: - `graph_pred`: Graph-level predictions (e.g., energy) of shape (n_graphs, graph_property_dim). - `stress_pred`: Stress predictions (if stress_head is provided) of shape (n_graphs, stress_dim). - `node_pred`: Node-level predictions of shape (n_nodes, node_property_dim). diff --git a/MindChemistry/mindchemistry/cell/orb/utils.py b/MindChemistry/mindchemistry/cell/orb/utils.py index 26797d9df..16eb538a4 100644 --- a/MindChemistry/mindchemistry/cell/orb/utils.py +++ b/MindChemistry/mindchemistry/cell/orb/utils.py @@ -240,7 +240,7 @@ def build_mlp( checkpoint: Whether to use checkpointing. Returns: - mlp: An MLP sequential container. + mlp. An MLP sequential container. """ # Size of each layer layer_sizes = [input_size] + hidden_layer_sizes diff --git a/docs/api_python/mindchemistry/cell/mindchemistry.cell.Allegro.rst b/docs/api_python/mindchemistry/cell/mindchemistry.cell.Allegro.rst index 117482316..0ca03b7c5 100644 --- a/docs/api_python/mindchemistry/cell/mindchemistry.cell.Allegro.rst +++ b/docs/api_python/mindchemistry/cell/mindchemistry.cell.Allegro.rst @@ -33,4 +33,16 @@ mindchemistry.cell.Allegro - **ValueError**: 如果 `new_tps_irreps` 的长度与 `tps_irreps` 不等。 - **ValueError**: 如果 `tps_irreps` 的阶数不为零。 - **ValueError**: 如果 `full_out_irreps` 的阶数不为零。 - - **ValueError**: 如果 `out_irreps` 的阶数不为零。 \ No newline at end of file + - **ValueError**: 如果 `out_irreps` 的阶数不为零。 + + .. py:method:: construct(embedding_data, edge_index, atom_types) + + 构建Allegro网络的计算图。 + + 参数： + - **embedding_data** (tuple(Tensor)) - 包含嵌入数据的元组，包括节点嵌入、边向量、边嵌入、边截断和边掩码等信息。 + - **edge_index** (Tensor) - 形状为 :math:`(2, edge\_num)` 的张量，表示边的连接关系，其中edge_index[1]表示中心节点，edge_index[0]表示邻居节点。 + - **atom_types** (Tensor) - 原子类型张量。 + + 返回： + Tensor，形状为 :math:`(edge\_num, final\_latent\_out)` 的边特征张量。 \ No newline at end of file diff --git a/docs/api_python/mindchemistry/cell/mindchemistry.cell.AutoEncoder.rst b/docs/api_python/mindchemistry/cell/mindchemistry.cell.AutoEncoder.rst index 01496a0f2..bdf8bd35a 100644 --- a/docs/api_python/mindchemistry/cell/mindchemistry.cell.AutoEncoder.rst +++ b/docs/api_python/mindchemistry/cell/mindchemistry.cell.AutoEncoder.rst @@ -25,3 +25,36 @@ mindchemistry.cell.AutoEncoder - **latents** (Tensor) - 形状为 :math:`(*, channels[-1])` 的张量。 - **x_recon** (Tensor) - 形状为 :math:`(*, channels[0])` 的张量。 + .. py:method:: construct(x) + + 构建自编码器网络的计算图。 + + 参数： + - **x** (Tensor) - 输入数据，形状为 :math:`(*, channels[0])` 的张量。 + + 返回： + tuple[Tensor]，包含两个张量： + + - **x_recon** (Tensor) - 重构的输入数据，形状为 :math:`(*, channels[0])` 的张量。 + - **latents** (Tensor) - 隐层编码，形状为 :math:`(*, channels[-1])` 的张量。 + + .. py:method:: encode(x) + + 将输入数据编码为隐层表示。 + + 参数： + - **x** (Tensor) - 输入数据，形状为 :math:`(*, channels[0])` 的张量。 + + 返回： + Tensor，隐层编码，形状为 :math:`(*, channels[-1])` 的张量。 + + .. py:method:: decode(z) + + 将隐层表示解码为重构数据。 + + 参数： + - **z** (Tensor) - 隐层编码，形状为 :math:`(*, channels[-1])` 的张量。 + + 返回： + Tensor，重构的输入数据，形状为 :math:`(*, channels[0])` 的张量。 + diff --git a/docs/api_python/mindchemistry/cell/mindchemistry.cell.FCNet.rst b/docs/api_python/mindchemistry/cell/mindchemistry.cell.FCNet.rst index 05bcc56ba..9f7823f68 100644 --- a/docs/api_python/mindchemistry/cell/mindchemistry.cell.FCNet.rst +++ b/docs/api_python/mindchemistry/cell/mindchemistry.cell.FCNet.rst @@ -23,3 +23,13 @@ mindchemistry.cell.FCNet 输出： - **output** (Tensor) - 形状为 :math:`(*, channels[-1])` 的Tensor。 + .. py:method:: construct(x) + + 构造函数，执行网络的前向传播。 + + 参数： + - **x** (Tensor) - 输入张量，形状为 :math:`(*, channels[0])`。 + + 返回： + Tensor，输出张量，形状为 :math:`(*, channels[-1])`。 + diff --git a/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.EnergyHead.rst b/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.EnergyHead.rst index fb549db32..001f02fcb 100644 --- a/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.EnergyHead.rst +++ b/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.EnergyHead.rst @@ -25,4 +25,19 @@ mindchemistry.cell.orb.EnergyHead 异常： - **ValueError** - 如果 `node_features` 中缺少必需的特征键。 - - **ValueError** - 如果 `node_aggregation` 不是支持的类型。 \ No newline at end of file + - **ValueError** - 如果 `node_aggregation` 不是支持的类型。 + + .. py:method:: predict(node_features, n_node, atom_types=None) + + 预测能量。 + + 参数： + - **node_features** (dict) - 节点特征字典，必须包含键"feat"，形状为 :math:`(n_{nodes}, latent\_dim)`。 + - **n_node** (Tensor) - 图中节点数量，形状为 :math:`(1,)`。 + - **atom_types** (Tensor, 可选) - 原子类型，形状为 :math:`(n_{nodes},)`。默认值： ``None``。 + + 返回： + dict，包含键"graph_pred"的字典，值的形状为 :math:`(1, target\_property\_dim)`。 + + 异常： + - **ValueError** - 如果 `node_features` 中缺少必需的特征键。 \ No newline at end of file diff --git a/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.GraphHead.rst b/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.GraphHead.rst index 75ae5ad7c..9fa44e0b4 100644 --- a/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.GraphHead.rst +++ b/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.GraphHead.rst @@ -22,4 +22,20 @@ mindchemistry.cell.orb.GraphHead - **output** (dict) - 包含键"stress_pred"的字典，值的形状为 :math:`(1, target\_property\_dim)`。 异常： - - **ValueError** - 如果 `node_features` 中缺少必需的特征键。 \ No newline at end of file + - **ValueError** - 如果 `node_features` 中缺少必需的特征键。 + + .. py:method:: predict(node_features, n_node, atomic_numbers=None) + + 预测图级属性。 + + 参数： + - **node_features** (dict) - 节点特征字典，必须包含键"feat"，形状为 :math:`(n_{nodes}, latent\_dim)`。 + - **n_node** (Tensor) - 图中节点数量，形状为 :math:`(1,)`。 + - **atomic_numbers** (Tensor, 可选) - 原子序数，用于参考能量计算。当compute_stress为False时必须提供。默认值： ``None``。 + + 返回： + Tensor，图级预测结果，形状为 :math:`(n_{graphs}, target\_property\_dim)`。如果compute_stress为True，则为应力张量；如果compute_stress为False，则为图级属性（例如能量）。 + + 异常： + - **AssertionError** - 如果compute_stress为False且未提供atomic_numbers。 + Tensor，图级预测结果，形状为 :math:`(n_{graphs}, target\_property\_dim)`。如果 `compute_stress` 为 ``True``，则为应力张量；如果 `compute_stress` 为 ``False``，则为图级属性（例如能量）。 \ No newline at end of file diff --git a/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.MoleculeGNS.rst b/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.MoleculeGNS.rst index c44551f32..7e702bd57 100644 --- a/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.MoleculeGNS.rst +++ b/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.MoleculeGNS.rst @@ -31,4 +31,24 @@ mindchemistry.cell.orb.MoleculeGNS 异常： - **ValueError** - 如果 `edge_features` 或 `node_features` 中缺少必需的特征键。 - - **ValueError** - 如果 `interactions` 不是支持的类型。 \ No newline at end of file + - **ValueError** - 如果 `interactions` 不是支持的类型。 + + .. py:method:: featurize_nodes(node_features) + + 特征化图的节点。 + + 参数： + - **node_features** (Dict[str, Tensor]) - 节点特征字典。 + + 返回： + Dict[str, Tensor]，带有原子嵌入的更新节点特征。 + + .. py:method:: featurize_edges(edge_features) + + 特征化图的边。 + + 参数： + - **edge_features** (Dict[str, Tensor]) - 边特征字典。 + + 返回： + Dict[str, Tensor]，带有径向基函数和单位向量的更新边特征。 \ No newline at end of file diff --git a/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.NodeHead.rst b/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.NodeHead.rst index 2e422d861..1d198d9c7 100644 --- a/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.NodeHead.rst +++ b/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.NodeHead.rst @@ -23,4 +23,15 @@ mindchemistry.cell.orb.NodeHead - **output** (dict) - 包含键 "node_pred" 的字典，值的形状为 :math:`(n_{nodes}, target\_property\_dim)`。 异常： - - **ValueError** - 如果 `node_features` 中缺少必需的特征键。 \ No newline at end of file + - **ValueError** - 如果 `node_features` 中缺少必需的特征键。 + + .. py:method:: predict(node_features, n_node) + + 预测节点级属性。 + + 参数： + - **node_features** (dict) - 节点特征字典，必须包含键 "feat"，形状为 :math:`(n_{nodes}, latent\_dim)`。 + - **n_node** (Tensor) - 图中节点数量，形状为 :math:`(1,)`。 + + 返回： + Tensor，节点级预测结果，形状为 :math:`(n_{nodes}, target\_property\_dim)`。 \ No newline at end of file diff --git a/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.Orb.rst b/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.Orb.rst index fe1d53a26..253508696 100644 --- a/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.Orb.rst +++ b/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.Orb.rst @@ -32,4 +32,23 @@ mindchemistry.cell.orb.Orb 异常： - **ValueError** - 如果既未提供node_head也未提供graph_head。 - **ValueError** - 如果cutoff_layers超过基础模型中的消息传递步骤数。 - - **ValueError** - 如果graph_head需要时未提供atomic_numbers。 \ No newline at end of file + - **ValueError** - 如果graph_head需要时未提供atomic_numbers。 + + .. py:method:: predict(edge_features, node_features, senders, receivers, n_node, atomic_numbers) + + 预测节点和/或图级属性。 + + 参数： + - **edge_features** (dict) - 边特征字典（例如，`{"vectors": Tensor, "r": Tensor}`）。 + - **node_features** (dict) - 节点特征字典（例如，`{"atomic_numbers": Tensor, "positions": Tensor, "atomic_numbers_embedding": Tensor}`）。 + - **senders** (Tensor) - 每条边的发送节点索引。形状：:math:`(n_{edges},)`。 + - **receivers** (Tensor) - 每条边的接收节点索引。形状：:math:`(n_{edges},)`。 + - **n_node** (Tensor) - 批次中每个图的节点数。形状：:math:`(n_{graphs},)`。 + - **atomic_numbers** (Tensor) - 用于参考能量计算的原子序数张量。 + + 返回： + dict，包含以下预测结果： + + - **graph_pred** (Tensor) - 图级预测（例如，能量）。形状：:math:`(n_{graphs}, graph\_property\_dim)`。 + - **stress_pred** (Tensor) - 应力预测（如果提供了stress_head）。形状：:math:`(n_{graphs}, stress\_dim)`。 + - **node_pred** (Tensor) - 节点级预测。形状：:math:`(n_{nodes}, node\_property\_dim)`。 \ No newline at end of file -- Gitee From 7e1f1d7b661c28f26ad8d4a4bc5446a92627e48a Mon Sep 17 00:00:00 2001 From: birfied Date: Mon, 15 Sep 2025 22:24:11 +0800 Subject: [PATCH 2/2] docs: correction --- .../mindchemistry/cell/mindchemistry.cell.orb.GraphHead.rst | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.GraphHead.rst b/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.GraphHead.rst index 9fa44e0b4..7c439e46e 100644 --- a/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.GraphHead.rst +++ b/docs/api_python/mindchemistry/cell/mindchemistry.cell.orb.GraphHead.rst @@ -37,5 +37,4 @@ mindchemistry.cell.orb.GraphHead Tensor，图级预测结果，形状为 :math:`(n_{graphs}, target\_property\_dim)`。如果compute_stress为True，则为应力张量；如果compute_stress为False，则为图级属性（例如能量）。 异常： - - **AssertionError** - 如果compute_stress为False且未提供atomic_numbers。 - Tensor，图级预测结果，形状为 :math:`(n_{graphs}, target\_property\_dim)`。如果 `compute_stress` 为 ``True``，则为应力张量；如果 `compute_stress` 为 ``False``，则为图级属性（例如能量）。 \ No newline at end of file + - **AssertionError** - 如果compute_stress为False且未提供atomic_numbers。 \ No newline at end of file -- Gitee