NepTrainKit 2.2.0__tar.gz → 2.2.2.dev23__tar.gz

{neptrainkit-2.2.0 → neptrainkit-2.2.2.dev23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NepTrainKit
-Version: 2.2.0
+Version: 2.2.2.dev23
 Summary: NepTrainKit is a Python package for visualizing and manipulating training datasets for NEP.
 Author: Chen Cheng bing
 Author-email: Chen Cheng bing <1747193328@qq.com>

{neptrainkit-2.2.0 → neptrainkit-2.2.2.dev23}/src/NepTrainKit/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.2.0'
-__version_tuple__ = version_tuple = (2, 2, 0)
+__version__ = version = '2.2.2.dev23'
+__version_tuple__ = version_tuple = (2, 2, 2, 'dev23')

{neptrainkit-2.2.0 → neptrainkit-2.2.2.dev23}/src/NepTrainKit/core/canvas/base/canvas.py

@@ -156,6 +156,16 @@ class CanvasLayoutBase(CanvasBase):
             self.nep_result_data.select(structure_index)
 
             self.update_scatter_color(structure_index, Brushes.Selected)
+
+    def inverse_select(self):
+        if self.nep_result_data is None:
+            return
+
+        active_indices = set(self.nep_result_data.structure.now_indices.tolist())
+        selected = set(self.nep_result_data.select_index)
+
+        self.select_index(list(selected), True)
+        self.select_index(list(active_indices - selected), False)
 class VispyCanvasLayoutBase(CanvasLayoutBase,QObject,metaclass=CombinedMeta):
     def __init__(self,*args,**kwargs):
         QObject.__init__(self)

{neptrainkit-2.2.0 → neptrainkit-2.2.2.dev23}/src/NepTrainKit/core/custom_widget/__init__.py

@@ -9,6 +9,8 @@ from .completer import CompleterModel, JoinDelegate, ConfigCompleter
 from .dialog import (
     GetIntMessageBox,
     SparseMessageBox,
+    IndexSelectMessageBox,
+    ShiftEnergyMessageBox,
     ProgressDialog,
     PeriodicTableDialog,
 )
@@ -19,6 +21,7 @@ from .card_widget import (
     MakeDataCardWidget,
 )
 from .doping_rule import DopingRulesWidget
+from .vacancy_rule import VacancyRulesWidget
 
 from .docker import MakeWorkflowArea
 from .search_widget import ConfigTypeSearchLineEdit
@@ -32,6 +35,8 @@ __all__ = [
     "ConfigCompleter",
     "GetIntMessageBox",
     "SparseMessageBox",
+    "IndexSelectMessageBox",
+    "ShiftEnergyMessageBox",
     "ProgressDialog",
     "PeriodicTableDialog",
     "SpinBoxUnitInputFrame",
@@ -43,4 +48,5 @@ __all__ = [
     "MyComboBoxSettingCard",
     "DoubleSpinBoxSettingCard",
     "DopingRulesWidget",
+    "VacancyRulesWidget",
 ]

{neptrainkit-2.2.0 → neptrainkit-2.2.2.dev23}/src/NepTrainKit/core/custom_widget/completer.py

@@ -33,7 +33,7 @@ class CompleterModel(QAbstractListModel):
     def data(self, index, role=Qt.ItemDataRole.DisplayRole):
 
         if not index.isValid():
-            print("index")
+
             return None
         # print(role)
         # Qt.ItemDataRole.DisplayRole
@@ -44,7 +44,7 @@ class CompleterModel(QAbstractListModel):
             if role == Qt.ItemDataRole.DisplayRole or role == Qt.ItemDataRole.EditRole:
                 return word
             elif role == CountRole:
-                print(count)
+
                 return str(count)
 
 

{neptrainkit-2.2.0 → neptrainkit-2.2.2.dev23}/src/NepTrainKit/core/custom_widget/dialog.py

@@ -4,14 +4,16 @@
 # @Author  : 兵
 # @email    : 1747193328@qq.com
 from PySide6.QtGui import QIcon
-from PySide6.QtWidgets import QVBoxLayout, QFrame, QGridLayout, QPushButton
+from PySide6.QtWidgets import QVBoxLayout, QFrame, QGridLayout, QPushButton, QLineEdit
 from PySide6.QtCore import Signal, Qt
 from qfluentwidgets import (
     MessageBoxBase,
     SpinBox,
     CaptionLabel,
     DoubleSpinBox,
+    CheckBox,
     ProgressBar,
+    ComboBox,
     FluentStyleSheet,
     FluentTitleBar,
     TitleLabel
@@ -73,6 +75,78 @@ class SparseMessageBox(MessageBoxBase):
 
         self.widget.setMinimumWidth(200)
 
+
+class IndexSelectMessageBox(MessageBoxBase):
+    """Dialog for selecting structures by index."""
+
+    def __init__(self, parent=None, tip="Specify index or slice"):
+        super().__init__(parent)
+        self.titleLabel = CaptionLabel(tip, self)
+        self.titleLabel.setWordWrap(True)
+        self.indexEdit = QLineEdit(self)
+        self.checkBox = CheckBox("Use original indices", self)
+        self.checkBox.setChecked(True)
+
+        self.viewLayout.addWidget(self.titleLabel)
+        self.viewLayout.addWidget(self.indexEdit)
+        self.viewLayout.addWidget(self.checkBox)
+
+        self.yesButton.setText('Ok')
+        self.cancelButton.setText('Cancel')
+        self.widget.setMinimumWidth(200)
+
+
+class ShiftEnergyMessageBox(MessageBoxBase):
+    """Dialog for energy baseline shift parameters."""
+
+    def __init__(self, parent=None, tip="Group regex patterns (comma separated)"):
+        super().__init__(parent)
+        self.titleLabel = CaptionLabel(tip, self)
+        self.titleLabel.setWordWrap(True)
+        self.groupEdit = QLineEdit(self)
+
+        self._frame = QFrame(self)
+        self.frame_layout = QGridLayout(self._frame)
+        self.frame_layout.setContentsMargins(0, 0, 0, 0)
+        self.frame_layout.setSpacing(2)
+
+        self.genSpinBox = SpinBox(self)
+        self.genSpinBox.setMaximum(100000000)
+        self.sizeSpinBox = SpinBox(self)
+        self.sizeSpinBox.setMaximum(999999)
+        self.tolSpinBox = DoubleSpinBox(self)
+        self.tolSpinBox.setDecimals(8)
+        self.tolSpinBox.setMinimum(0)
+        self.modeCombo = ComboBox(self)
+        self.modeCombo.addItems([
+            "REF_GROUP_ALIGNMENT",
+            "ZERO_BASELINE_ALIGNMENT",
+            "DFT_TO_NEP_ALIGNMENT",
+        ])
+        self.modeCombo.setCurrentText("DFT_TO_NEP_ALIGNMENT")
+
+
+        self.frame_layout.addWidget(CaptionLabel("Max generations", self), 0, 0)
+        self.frame_layout.addWidget(self.genSpinBox, 0, 1)
+        self.frame_layout.addWidget(CaptionLabel("Population size", self), 1, 0)
+        self.frame_layout.addWidget(self.sizeSpinBox, 1, 1)
+        self.frame_layout.addWidget(CaptionLabel("Convergence tol", self), 2, 0)
+        self.frame_layout.addWidget(self.tolSpinBox, 2, 1)
+        self.frame_layout.addWidget(CaptionLabel("Mode", self), 3, 0)
+        self.frame_layout.addWidget(self.modeCombo, 3, 1)
+
+
+        self.viewLayout.addWidget(self.titleLabel)
+        self.viewLayout.addWidget(self.groupEdit)
+        self.viewLayout.addWidget(self._frame)
+
+        self.yesButton.setText('Ok')
+        self.cancelButton.setText('Cancel')
+        self.widget.setMinimumWidth(250)
+
+
+
+
 class ProgressDialog(FramelessDialog):
     """进度条弹窗"""
     def __init__(self,parent=None,title=""):

{neptrainkit-2.2.0 → neptrainkit-2.2.2.dev23}/src/NepTrainKit/core/custom_widget/doping_rule.py

@@ -15,20 +15,19 @@ from PySide6.QtWidgets import (
     QGridLayout,
     QHBoxLayout,
     QVBoxLayout,
-    QWidget, QLineEdit,
-
+    QWidget,
+    QLineEdit,
 )
 from qfluentwidgets import (
     BodyLabel,
     TransparentToolButton,
-    SpinBox,
-    DoubleSpinBox,
     FluentIcon,
     LineEdit,
     RadioButton,
     ToolTipFilter,
     ToolTipPosition,
 )
+from .input import SpinBoxUnitInputFrame
 
 
 class DopingRuleItem(QFrame):
@@ -43,16 +42,20 @@ class DopingRuleItem(QFrame):
         self.target_edit = QLineEdit(self)
         self.target_edit.setPlaceholderText("Cs")
 
-        self.setFixedSize(300, 100)
+        self.setFixedSize(300, 130)
         self.dopants_edit = QLineEdit(self)
          
-        self.concentration_spin = QLineEdit(self)
-        self.concentration_spin.setText("1")
+        self.concentration_frame = SpinBoxUnitInputFrame(self)
+        self.concentration_frame.set_input(["-", ""], 2, "float")
+        self.concentration_frame.setRange(0, 1)
+        self.concentration_frame.set_input_value([1.0, 1.0])
 
         self.concentration_botton = RadioButton("Conc", self)
         self.concentration_botton.setChecked(True)
-        self.count_spin = QLineEdit(self)
-        self.count_spin.setText("10")
+        self.count_frame = SpinBoxUnitInputFrame(self)
+        self.count_frame.set_input(["-", ""], 2, "int")
+        self.count_frame.setRange(0, 10000)
+        self.count_frame.set_input_value([10, 10])
         self.count_botton = RadioButton("Count", self)
 
         self.indices_edit = QLineEdit(self)
@@ -78,15 +81,15 @@ class DopingRuleItem(QFrame):
         self.concentration_botton.setToolTip("Use concentration")
         self.concentration_botton.installEventFilter(ToolTipFilter(self.concentration_botton, 300, ToolTipPosition.TOP))
         self.layout.addWidget(self.concentration_botton, 2, 0)
-        self.layout.addWidget(self.concentration_spin, 2, 1)
+        self.layout.addWidget(self.concentration_frame, 2, 1,1,4)
         self.count_botton.setToolTip("Use count")
         self.count_botton.installEventFilter(ToolTipFilter(self.count_botton, 300, ToolTipPosition.TOP))
-        self.layout.addWidget(self.count_botton, 2, 2)
-        self.layout.addWidget(self.count_spin, 2, 3)
+        self.layout.addWidget(self.count_botton, 3, 0)
+        self.layout.addWidget(self.count_frame, 3, 1,1,4)
 
         self.delete_button.setToolTip("Delete rule")
         self.delete_button.installEventFilter(ToolTipFilter(self.delete_button, 300, ToolTipPosition.TOP))
-        self.layout.addWidget(self.delete_button, 0, 4, 3, 3)
+        self.layout.addWidget(self.delete_button, 0, 4, 3, 1)
 
     def _delete_self(self) -> None:
         self.setParent(None)
@@ -115,9 +118,9 @@ class DopingRuleItem(QFrame):
         except Exception:
             logger.error(traceback.format_exc())
 
-        rule["concentration"] =float(self.concentration_spin.text().strip()) if self.concentration_spin.text().strip() else 0
+        rule["concentration"] = [float(v) for v in self.concentration_frame.get_input_value()]
 
-        rule["count"] =float(self.count_spin.text().strip()) if self.count_spin.text().strip() else 0
+        rule["count"] = [int(v) for v in self.count_frame.get_input_value()]
         rule["use"] = "concentration" if self.concentration_botton.isChecked() else "count"
         indices_text = self.indices_edit.text().strip()
         if indices_text:
@@ -136,9 +139,9 @@ class DopingRuleItem(QFrame):
         if dopants is not None:
             self.dopants_edit.setText(json.dumps(dopants))
         if "concentration" in rule:
-            self.concentration_spin.setText(str(rule["concentration"]))
+            self.concentration_frame.set_input_value(rule["concentration"])
         if "count" in rule:
-            self.count_spin.setText(str(rule["count"]))
+            self.count_frame.set_input_value(rule["count"])
         if "group" in rule:
             self.indices_edit.setText(",".join(str(i) for i in rule["group"]))
         if "use" in  rule:

{neptrainkit-2.2.0 → neptrainkit-2.2.2.dev23}/src/NepTrainKit/core/custom_widget/input.py

@@ -53,5 +53,8 @@ class SpinBoxUnitInputFrame(QFrame):
         return [input_object.value() for input_object in self.object_list]
 
     def set_input_value(self, value_list):
+        if not isinstance(value_list,list):
+            value_list=[value_list]*len(self.object_list)
+
         for i, input_object in enumerate(self.object_list):
             input_object.setValue(value_list[i])

neptrainkit-2.2.2.dev23/src/NepTrainKit/core/custom_widget/vacancy_rule.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""Widget to edit vacancy rules."""
+
+from __future__ import annotations
+
+from PySide6.QtCore import Qt
+from PySide6.QtGui import QIcon
+from PySide6.QtWidgets import (
+    QFrame,
+    QGridLayout,
+    QHBoxLayout,
+    QVBoxLayout,
+    QWidget,
+    QLineEdit,
+)
+from qfluentwidgets import (
+    BodyLabel,
+    TransparentToolButton,
+    FluentIcon,
+    ToolTipFilter,
+    ToolTipPosition,
+)
+from .input import SpinBoxUnitInputFrame
+
+
+class VacancyRuleItem(QFrame):
+    """Single vacancy rule widget."""
+
+    def __init__(self, parent: QWidget | None = None) -> None:
+        super().__init__(parent)
+        self.layout = QGridLayout(self)
+        self.layout.setContentsMargins(0, 0, 0, 0)
+        self.layout.setSpacing(4)
+        self.setStyleSheet("background-color: rgb(239, 249, 254);")
+
+        self.element_edit = QLineEdit(self)
+        self.element_edit.setPlaceholderText("Cs")
+        self.group_edit = QLineEdit(self)
+        self.count_frame = SpinBoxUnitInputFrame(self)
+        self.count_frame.set_input(["-", ""], 2, "int")
+        self.count_frame.setRange(0, 10000)
+        self.count_frame.set_input_value([1, 1])
+
+        self.delete_button = TransparentToolButton(QIcon(":/images/src/images/delete.svg"), self)
+        self.delete_button.clicked.connect(self._delete_self)
+
+        self.element_label = BodyLabel("Element", self)
+        self.element_label.setToolTip("Element to remove")
+        self.element_label.installEventFilter(ToolTipFilter(self.element_label, 300, ToolTipPosition.TOP))
+        self.group_label = BodyLabel("Group", self)
+        self.group_label.setToolTip("Optional group name")
+        self.group_label.installEventFilter(ToolTipFilter(self.group_label, 300, ToolTipPosition.TOP))
+        self.count_label = BodyLabel("Count", self)
+        self.count_label.setToolTip("Number of atoms to remove")
+        self.count_label.installEventFilter(ToolTipFilter(self.count_label, 300, ToolTipPosition.TOP))
+
+        self.layout.addWidget(self.element_label, 0, 0)
+        self.layout.addWidget(self.element_edit, 0, 1)
+        self.layout.addWidget(self.group_label, 0, 2)
+        self.layout.addWidget(self.group_edit, 0, 3)
+        self.layout.addWidget(self.count_label, 1, 0)
+        self.layout.addWidget(self.count_frame, 1, 1)
+        self.layout.addWidget(self.delete_button, 0, 4, 2, 1)
+
+    def _delete_self(self) -> None:
+        self.setParent(None)
+        self.deleteLater()
+
+    def to_rule(self) -> dict:
+        rule: dict[str, object] = {}
+        element = self.element_edit.text().strip()
+        if element:
+            rule["element"] = element
+        rule["count"] = [int(v) for v in self.count_frame.get_input_value()]
+        groups = self.group_edit.text().strip()
+        if groups:
+            rule["group"] = [g.strip() for g in groups.split(",") if g.strip()]
+        return rule
+
+    def from_rule(self, rule: dict) -> None:
+        if not rule:
+            return
+        self.element_edit.setText(str(rule.get("element", "")))
+        if "count" in rule:
+            self.count_frame.set_input_value(rule["count"])
+        if "group" in rule:
+            self.group_edit.setText(",".join(str(i) for i in rule["group"]))
+
+
+class VacancyRulesWidget(QWidget):
+    """Container widget for multiple vacancy rules."""
+
+    def __init__(self, parent: QWidget | None = None) -> None:
+        super().__init__(parent)
+        self.layout = QVBoxLayout(self)
+        self.layout.setContentsMargins(0, 0, 0, 0)
+        self.layout.setSpacing(4)
+
+        btn_layout = QHBoxLayout()
+        btn_layout.setContentsMargins(0, 0, 0, 0)
+        self.add_button = TransparentToolButton(FluentIcon.ADD, self)
+        self.add_button.clicked.connect(self.add_rule)
+        self.add_button.setToolTip("Add rule")
+        self.add_button.installEventFilter(ToolTipFilter(self.add_button, 300, ToolTipPosition.TOP))
+        btn_layout.addWidget(self.add_button, 0, Qt.AlignLeft)
+        btn_layout.addStretch(1)
+        self.layout.addLayout(btn_layout)
+
+        self.rule_container = QWidget(self)
+        self.rule_layout = QVBoxLayout(self.rule_container)
+        self.rule_layout.setContentsMargins(0, 0, 0, 0)
+        self.rule_layout.setSpacing(4)
+        self.layout.addWidget(self.rule_container)
+
+    def add_rule(self, rule: dict | None = None) -> VacancyRuleItem:
+        item = VacancyRuleItem(self.rule_container)
+        self.rule_layout.addWidget(item)
+        if rule:
+            item.from_rule(rule)
+        return item
+
+    def to_rules(self) -> list[dict]:
+        rules: list[dict] = []
+        for i in range(self.rule_layout.count()):
+            widget = self.rule_layout.itemAt(i).widget()
+            if isinstance(widget, VacancyRuleItem):
+                rule = widget.to_rule()
+                if rule:
+                    rules.append(rule)
+        return rules
+
+    def from_rules(self, rules: list[dict]) -> None:
+        while self.rule_layout.count():
+            item = self.rule_layout.takeAt(0).widget()
+            if item is not None:
+                item.deleteLater()
+        for rule in rules or []:
+            self.add_rule(rule)

neptrainkit-2.2.2.dev23/src/NepTrainKit/core/energy_shift.py

@@ -0,0 +1,216 @@
+"""Utilities for shifting structure energies using atomic baselines.
+抄的陈博的代码(已允许)
+url: https://github.com/brucefan1983/GPUMD/tree/master/tools/Analysis_and_Processing/energy-reference-aligner
+Zherui Chen Email: <chenzherui0124@foxmail.com>
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from collections import Counter
+from typing import List, Dict
+import re
+from NepTrainKit import utils
+from .structure import Structure
+from .calculator import NepCalculator
+
+REF_GROUP_ALIGNMENT = "REF_GROUP_ALIGNMENT"
+ZERO_BASELINE_ALIGNMENT = "ZERO_BASELINE_ALIGNMENT"
+DFT_TO_NEP_ALIGNMENT = "DFT_TO_NEP_ALIGNMENT"
+
+def longest_common_prefix(strs: List[str]) -> str:
+    if not strs:
+        return ""
+    s1, s2 = min(strs), max(strs)
+    for i, c in enumerate(s1):
+        if c != s2[i]:
+            return s1[:i]
+    return s1
+
+
+def suggest_group_patterns(config_types: List[str], min_group_size: int = 2, min_prefix_len: int = 3) -> List[str]:
+    """Group strings by common prefix without relying on delimiters, and output regex patterns."""
+    unused = set(config_types)
+    patterns = []
+
+    while unused:
+        base = unused.pop()
+        group = [base]
+        to_remove = []
+
+        for other in unused:
+            prefix = longest_common_prefix([base, other])
+            if len(prefix) >= min_prefix_len:
+                group.append(other)
+                to_remove.append(other)
+
+        for item in to_remove:
+            unused.remove(item)
+
+        if len(group) >= min_group_size:
+            prefix = longest_common_prefix(group)
+            patterns.append(re.escape(prefix) + '.*')
+        else:
+            patterns.extend(re.escape(g) for g in group)
+
+    return sorted(patterns)
+def atomic_baseline_cost(param_population: np.ndarray,
+                         energies: np.ndarray,
+                         element_counts: np.ndarray,
+                         target_energies: np.ndarray) -> np.ndarray:
+    """Vectorized MSE cost for atomic reference baseline."""
+    shifted = energies[None, :] - np.dot(param_population, element_counts.T)
+    cost = np.mean((shifted - target_energies[None, :]) ** 2, axis=1)
+    return cost.reshape(-1, 1)
+
+@utils.timeit
+def nes_optimize_atomic_baseline(num_variables: int,
+                                 max_generations: int,
+                                 energies: np.ndarray,
+                                 element_counts: np.ndarray,
+                                 targets: np.ndarray,
+                                 pop_size: int = 40,
+                                 tol: float = 1e-8,
+                                 seed: int = 42,
+                                 print_every: int = 100) -> np.ndarray:
+    """NES optimizer for atomic reference energies."""
+    np.random.seed(seed)
+
+    best_fitness = np.ones((max_generations, 1))
+    elite = np.zeros((max_generations, num_variables))
+    mean = -1 * np.random.rand(1, num_variables)
+    stddev = 0.1 * np.ones((1, num_variables))
+    lr_mean = 1.0
+    lr_std = (3 + np.log(num_variables)) / (5 * np.sqrt(num_variables)) / 2
+    weights = np.maximum(0, np.log(pop_size / 2 + 1) - np.log(np.arange(1, pop_size + 1)))
+    weights = weights / np.sum(weights) - 1 / pop_size
+
+    for gen in range(max_generations):
+        z = np.random.randn(pop_size, num_variables)
+        pop = mean + stddev * z
+        fitness = atomic_baseline_cost(pop, energies, element_counts, targets)
+        idx = np.argsort(fitness.flatten())
+        fitness = fitness[idx]
+        z = z[idx, :]
+        pop = pop[idx, :]
+        best_fitness[gen] = fitness[0]
+        elite[gen, :] = pop[0, :]
+        mean += lr_mean * stddev * (weights @ z)
+        stddev *= np.exp(lr_std * (weights @ (z ** 2 - 1)))
+        if gen > 0 and abs(best_fitness[gen] - best_fitness[gen - 1]) < tol:
+            best_fitness = best_fitness[:gen + 1]
+            elite = elite[:gen + 1]
+            break
+    return elite[-1]
+
+
+
+def shift_dataset_energy(
+        structures: List[Structure],
+        reference_structures: List[Structure] | None,
+        max_generations: int = 100000,
+        population_size: int = 40,
+        convergence_tol: float = 1e-8,
+        random_seed: int = 42,
+        group_patterns: List[str] | None = None,
+        alignment_mode: str = REF_GROUP_ALIGNMENT,
+        nep_energy_array: np.array | None = None):
+    """Shift structure energies using different alignment strategies.
+
+    Parameters
+    ----------
+    structures
+        Structures whose energies will be shifted.
+    reference_structures
+        Structures used to compute the reference mean energy when
+        ``alignment_mode`` is ``REF_GROUP_ALIGNMENT``.
+    alignment_mode
+        One of ``REF_GROUP_ALIGNMENT``, ``ZERO_BASELINE_ALIGNMENT`` or
+        ``DFT_TO_NEP_ALIGNMENT``.
+    nep_energy_array
+        nep energy array when ``alignment_mode`` is
+        ``DFT_TO_NEP_ALIGNMENT``.
+    """
+    frames = []
+    for s in structures:
+        energy = float(s.energy)
+        config_type = str(s.additional_fields.get("Config_type", ""))
+        elem_counts = Counter(s.elements)
+
+        frames.append({"energy": energy, "config_type": config_type, "elem_counts": elem_counts})
+
+    all_elements = sorted({e for f in frames for e in f["elem_counts"]})
+    num_elements = len(all_elements)
+
+
+    ref_mean = None
+    if alignment_mode == REF_GROUP_ALIGNMENT:
+        if not len(reference_structures):
+            raise ValueError("reference_structures is required for REF_GROUP_ALIGNMENT")
+        ref_energies = np.array([f.energy for f in reference_structures])
+        ref_mean = np.mean(ref_energies)
+
+    if alignment_mode == DFT_TO_NEP_ALIGNMENT:
+        if nep_energy_array is None:
+            raise ValueError("nep_energy_array is required for DFT_TO_NEP_ALIGNMENT")
+
+        for f, e in zip(frames, nep_energy_array):
+            f["nep_energy"] = e * f["elem_counts"].total()
+
+    all_config_types = {f["config_type"] for f in frames}
+
+    # build mapping from config_type to regex group name
+    config_to_group: Dict[str, str] = {}
+    if group_patterns:
+        for pat in group_patterns:
+            try:
+                regex = re.compile(pat)
+            except re.error:
+                continue
+            for ct in all_config_types:
+                if ct not in config_to_group and regex.match(ct):
+                    config_to_group[ct] = pat
+    for ct in all_config_types:
+        config_to_group.setdefault(ct, ct)
+
+    shift_groups = sorted(set(config_to_group.values()))
+
+    group_to_atomic_ref = {}
+    for group in shift_groups:
+
+        grp_frames = [f for f in frames if config_to_group[f["config_type"]] == group]
+        if not grp_frames:
+            continue
+        energies = np.array([f["energy"] for f in grp_frames])
+        counts = np.array([[f["elem_counts"].get(e, 0) for e in all_elements] for f in grp_frames], dtype=float)
+
+        if alignment_mode == REF_GROUP_ALIGNMENT:
+            targets = np.full_like(energies, ref_mean)
+        elif alignment_mode == ZERO_BASELINE_ALIGNMENT:
+            targets = np.zeros_like(energies)
+        else:  # DFT_TO_NEP_ALIGNMENT
+            targets = np.array([f["nep_energy"] for f in grp_frames])
+        atomic_ref = nes_optimize_atomic_baseline(
+            num_elements,
+            max_generations,
+            energies,
+            counts,
+            targets,
+            pop_size=population_size,
+            tol=convergence_tol,
+            seed=random_seed,
+            print_every=100,
+        )
+        group_to_atomic_ref[group] = atomic_ref
+        #这里是为了更新ui信号
+        yield 1
+
+    # apply shift
+    for s, frame in zip(structures, frames):
+        group = config_to_group[frame["config_type"]]
+        if group in group_to_atomic_ref:
+            count_vec = np.array([frame["elem_counts"].get(e, 0) for e in all_elements], dtype=float)
+            shift = np.dot(count_vec, group_to_atomic_ref[group])
+            new_energy = frame["energy"] - shift
+            s.energy = new_energy
+    # return group_to_atomic_ref

{neptrainkit-2.2.0 → neptrainkit-2.2.2.dev23}/src/NepTrainKit/core/io/base.py

@@ -111,6 +111,9 @@ class NepData:
         """
         return self.data.now_data
     @property
+    def now_indices(self):
+        return self.data.now_indices
+    @property
     def all_data(self):
         return self.data.all_data
 

{neptrainkit-2.2.0 → neptrainkit-2.2.2.dev23}/src/NepTrainKit/core/io/nep.py

@@ -155,6 +155,17 @@ class ResultData(QObject):
                 self.select_index.remove(i)
 
         self.updateInfoSignal.emit()
+
+    def inverse_select(self):
+        """Invert the current selection state of all active structures"""
+        active_indices = set(self.structure.data.now_indices.tolist())
+        selected_indices = set(self.select_index)
+        unselect = list(selected_indices)
+        select = list(active_indices - selected_indices)
+        if unselect:
+            self.uncheck(unselect)
+        if select:
+            self.select(select)
     def export_selected_xyz(self,save_file_path):
         """
         导出当前选中的结构
@@ -166,7 +177,7 @@ class ResultData(QObject):
 
                 index=self.structure.convert_index(index)
 
-                for structure in self.structure.now_data[index]:
+                for structure in self.structure.all_data[index]:
                     structure.write(f)
 
             MessageManager.send_info_message(f"File exported to: {save_file_path}")