diff --git a/electronic_structure/configs/deeph/README.md b/electronic_structure/configs/deeph/README.md new file mode 100644 index 00000000..1b3435d6 --- /dev/null +++ b/electronic_structure/configs/deeph/README.md @@ -0,0 +1,198 @@ +# DeepH + +[Deep-learning density functional theory Hamiltonian for efficient ab initio electronic-structure calculation](https://www.nature.com/articles/s43588-022-00265-6) + +## Introduction + +DeepH learns density functional theory (DFT) Hamiltonian matrix elements from +crystal structures. By combining locality of electronic interactions with graph +neural network message passing, the model predicts Hamiltonian blocks on +periodic atom-pair edges and enables efficient downstream electronic-structure +calculation without repeatedly running expensive self-consistent DFT. + +![DeepH Fig. 2: Crystal graph and MPNN including L layers employed by DeepH](../../docs/deeph_fig2.png) + +Figure source: [Nature Computational Science, Fig. 2](https://www.nature.com/articles/s43588-022-00265-6/figures/2). + +## Model + +DeepH represents a crystal as a periodic graph. Nodes are atoms, edges are atom +pairs within a cutoff radius, and each edge corresponds to a Hamiltonian block +$H_{ij}$ between local atomic orbitals. The PaddleMaterials integration reuses +the suite graph converter to build the periodic crystal graph and keeps the +DeepH-specific LCMP subgraph features for angular message passing. + +For an edge $(i, j)$, the target is the selected orbital entry or block of +$H_{ij}$. Edge distances are expanded by Gaussian radial bases, angular +features are generated from spherical harmonics in local frames, and stacked +message-passing layers update atom/edge representations. A final edge decoder +predicts Hamiltonian entries with a masked mean-squared-error objective: + +$$ +\mathcal{L} = +\frac{\sum_{(i,j),k} M_{ij,k}(\hat{H}_{ij,k}-H_{ij,k})^2} +{\sum_{(i,j),k} M_{ij,k}}. +$$ + +The training config validates the non-spin graphene Hamiltonian path with +`interface=npz`, `target=hamiltonian`, and `if_lcmp_graph=True`. Both +radius-based graphs (`create_from_DFT=False`) and DFT matrix sparsity graphs +(`create_from_DFT=True`) are supported. Label-free inference uses +`interface=npz_rc_only` with local coordinates generated from OpenMX overlap +blocks. + +## Datasets + +The graphene example uses DeepH processed data. Each structure folder contains +the crystal geometry and Hamiltonian labels: + +```text +lat.dat +site_positions.dat +element.dat +orbital_types.dat +rh.npz +rc.npz +``` + +The PaddleMaterials dataset adapter: + +- reads DeepH processed folders from `raw_dir`; +- reconstructs canonical `pymatgen.Structure` objects with `BuildStructure`; +- builds periodic crystal graphs with `FindPointsInSpheres`; +- attaches DeepH Hamiltonian labels and LCMP subgraph metadata; +- caches parsed graph samples under the configured `graph_dir`. + +The dataset was released with the +[official DeepH-pack project](https://github.com/mzjb/DeepH-pack). Its processed +format stores crystal geometry, local-coordinate matrices, and DFT Hamiltonian +blocks for each structure. The train/validation/test split is generated with a +fixed random seed from the ratios below. + +The baseline graphene split follows the original DeepH config ratios: + +| Dataset | Train | Val | Test | Seed | +| --- | ---: | ---: | ---: | ---: | +| graphene | 60% | 20% | 20% | 42 | + +Download the graphene dataset, pretrained checkpoint, and logs from +[deeph_bce_handoff](https://pan.baidu.com/s/1dkXlzYV_3DLQugoQDKjsbA?pwd=vi6r) +(password: `vi6r`). + +## Results + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
Model NameDatasetForward max diffTrain align diffTest lossCompiler speedupConfigCheckpoint | Log
deeph_graphenegraphene1.1444e-054.0531e-060.0107244188.9434%deeph_graphenedeeph_bce_handoff (pwd: vi6r)
deeph_tbg_subsetTBG_subset1.3351e-051.6809e-05---alignment evidence
+ +**Metric notes:** `Train align diff` reports the step-1 loss difference against +the original DeepH reference implementation. The graphene supervised metric +alignment has absolute `test_loss` difference `1.0801e-04`. Compiler speedup is +measured by comparing Paddle dynamic evaluation latency (`21.6952 ms`) with +`to_static`/CINN evaluation latency (`11.4824 ms`). + +## Command + +### Training + +```bash +python electronic_structure/train.py \ + -c electronic_structure/configs/deeph/deeph_graphene.yaml +``` + +### Validation + +```bash +python electronic_structure/train.py \ + -c electronic_structure/configs/deeph/deeph_graphene.yaml \ + Global.do_train=False Global.do_eval=True Global.do_test=False \ + Trainer.pretrained_model_path='data/deeph/graphene/best.pdparams' +``` + +### Testing + +```bash +python electronic_structure/train.py \ + -c electronic_structure/configs/deeph/deeph_graphene.yaml \ + Global.do_train=False Global.do_eval=False Global.do_test=True \ + Trainer.pretrained_model_path='data/deeph/graphene/best.pdparams' +``` + +### Dynamic-to-static / CINN evaluation + +```bash +FLAGS_use_cinn=1 python electronic_structure/train.py \ + -c electronic_structure/configs/deeph/deeph_graphene.yaml \ + Global.do_train=False Global.do_eval=True Global.do_test=False \ + Trainer.pretrained_model_path='data/deeph/graphene/best.pdparams' +``` + +### OpenMX DFT inference + +DeepH inference first runs the overlap-only OpenMX calculation through ASE, +parses its sparse overlap blocks, constructs the LCMP graph, runs the Paddle +checkpoint, and writes the predicted Hamiltonian. Configure the external DFT +software before launching ppmatSim: + +```bash +export OPENMX_DFT_DATA_PATH=/path/to/DFT_DATA19 + +python ppmatSim/main.py --config-name deeph_openmx \ + Model.config_path=electronic_structure/configs/deeph/deeph_graphene.yaml \ + Model.checkpoint_path=data/deeph/graphene/best.pdparams \ + Calculator.command='mpirun -np 4 /path/to/openmx_overlap' \ + System.file_path=/path/to/inference_cifs +``` + +For each input structure, the workflow saves the raw OpenMX calculation, +processed geometry, `overlaps.npz`, `rc.npz`, and the predicted `rh_pred.npz` +under `deeph_dft_results//`. The supplied graphene model supports +carbon structures with the OpenMX `C6.0-s2p2d1` basis. + +## Citation + +```bibtex +@article{li2022deeph, + title={Deep-learning density functional theory Hamiltonian for efficient ab initio electronic-structure calculation}, + author={Li, He and Wang, Zun and Zou, Nianlong and Ye, Meng and Xu, Runzhang and Gong, Xiaoxun and Duan, Wenhui and Xu, Yong}, + journal={Nature Computational Science}, + volume={2}, + number={6}, + pages={367--377}, + year={2022}, + publisher={Nature Publishing Group} +} +``` diff --git a/electronic_structure/configs/deeph/deeph_graphene.yaml b/electronic_structure/configs/deeph/deeph_graphene.yaml new file mode 100644 index 00000000..91710680 --- /dev/null +++ b/electronic_structure/configs/deeph/deeph_graphene.yaml @@ -0,0 +1,401 @@ +Global: + do_train: true + do_eval: true + do_test: true +Predict: + eval_with_no_grad: true +Trainer: + max_epochs: 5 + seed: 42 + output_dir: ./output/deeph_graphene + save_freq: 1 + log_freq: 1 + start_eval_epoch: 0 + eval_freq: 1 + pretrained_model_path: null + pretrained_weight_name: null + resume_from_checkpoint: null + use_amp: false + amp_level: O1 + eval_with_no_grad: true + gradient_accumulation_steps: 1 + best_metric_indicator: eval_loss + name_for_best_metric: loss + greater_is_better: false + compute_metric_during_train: false + metric_strategy_during_eval: epoch + use_visualdl: false + use_wandb: false + use_tensorboard: false +Model: + __class_name__: DeepHHamiltonian + __init_params__: + num_species: 1 + in_atom_fea_len: 64 + in_edge_fea_len: 128 + num_orbital: 169 + num_l: 5 + gauss_stop: 6.0 + if_exp: true + normalization: LayerNorm + target_name: label +Optimizer: + __class_name__: Adam + __init_params__: + lr: 0.001 + weight_decay: 0.0 + beta1: 0.9 + beta2: 0.999 + epsilon: 1.0e-08 + clip_norm_global: 4.2 +DeepH: + basic: + graph_dir: ./data/deeph/graphene/graph_data + save_dir: ./output/deeph_graphene + raw_dir: ./data/deeph/graphene/processed + dataset_name: graphene_local_gpu + only_get_graph: false + interface: npz + target: hamiltonian + disable_cuda: false + device: cpu + num_threads: 8 + save_to_time_folder: false + save_csv: false + tb_writer: false + seed: 42 + multiprocessing: 0 + orbital: + - 6 6: [0, 0] + - 6 6: [0, 1] + - 6 6: [0, 2] + - 6 6: [0, 3] + - 6 6: [0, 4] + - 6 6: [0, 5] + - 6 6: [0, 6] + - 6 6: [0, 7] + - 6 6: [0, 8] + - 6 6: [0, 9] + - 6 6: [0, 10] + - 6 6: [0, 11] + - 6 6: [0, 12] + - 6 6: [1, 0] + - 6 6: [1, 1] + - 6 6: [1, 2] + - 6 6: [1, 3] + - 6 6: [1, 4] + - 6 6: [1, 5] + - 6 6: [1, 6] + - 6 6: [1, 7] + - 6 6: [1, 8] + - 6 6: [1, 9] + - 6 6: [1, 10] + - 6 6: [1, 11] + - 6 6: [1, 12] + - 6 6: [2, 0] + - 6 6: [2, 1] + - 6 6: [2, 2] + - 6 6: [2, 3] + - 6 6: [2, 4] + - 6 6: [2, 5] + - 6 6: [2, 6] + - 6 6: [2, 7] + - 6 6: [2, 8] + - 6 6: [2, 9] + - 6 6: [2, 10] + - 6 6: [2, 11] + - 6 6: [2, 12] + - 6 6: [3, 0] + - 6 6: [3, 1] + - 6 6: [3, 2] + - 6 6: [3, 3] + - 6 6: [3, 4] + - 6 6: [3, 5] + - 6 6: [3, 6] + - 6 6: [3, 7] + - 6 6: [3, 8] + - 6 6: [3, 9] + - 6 6: [3, 10] + - 6 6: [3, 11] + - 6 6: [3, 12] + - 6 6: [4, 0] + - 6 6: [4, 1] + - 6 6: [4, 2] + - 6 6: [4, 3] + - 6 6: [4, 4] + - 6 6: [4, 5] + - 6 6: [4, 6] + - 6 6: [4, 7] + - 6 6: [4, 8] + - 6 6: [4, 9] + - 6 6: [4, 10] + - 6 6: [4, 11] + - 6 6: [4, 12] + - 6 6: [5, 0] + - 6 6: [5, 1] + - 6 6: [5, 2] + - 6 6: [5, 3] + - 6 6: [5, 4] + - 6 6: [5, 5] + - 6 6: [5, 6] + - 6 6: [5, 7] + - 6 6: [5, 8] + - 6 6: [5, 9] + - 6 6: [5, 10] + - 6 6: [5, 11] + - 6 6: [5, 12] + - 6 6: [6, 0] + - 6 6: [6, 1] + - 6 6: [6, 2] + - 6 6: [6, 3] + - 6 6: [6, 4] + - 6 6: [6, 5] + - 6 6: [6, 6] + - 6 6: [6, 7] + - 6 6: [6, 8] + - 6 6: [6, 9] + - 6 6: [6, 10] + - 6 6: [6, 11] + - 6 6: [6, 12] + - 6 6: [7, 0] + - 6 6: [7, 1] + - 6 6: [7, 2] + - 6 6: [7, 3] + - 6 6: [7, 4] + - 6 6: [7, 5] + - 6 6: [7, 6] + - 6 6: [7, 7] + - 6 6: [7, 8] + - 6 6: [7, 9] + - 6 6: [7, 10] + - 6 6: [7, 11] + - 6 6: [7, 12] + - 6 6: [8, 0] + - 6 6: [8, 1] + - 6 6: [8, 2] + - 6 6: [8, 3] + - 6 6: [8, 4] + - 6 6: [8, 5] + - 6 6: [8, 6] + - 6 6: [8, 7] + - 6 6: [8, 8] + - 6 6: [8, 9] + - 6 6: [8, 10] + - 6 6: [8, 11] + - 6 6: [8, 12] + - 6 6: [9, 0] + - 6 6: [9, 1] + - 6 6: [9, 2] + - 6 6: [9, 3] + - 6 6: [9, 4] + - 6 6: [9, 5] + - 6 6: [9, 6] + - 6 6: [9, 7] + - 6 6: [9, 8] + - 6 6: [9, 9] + - 6 6: [9, 10] + - 6 6: [9, 11] + - 6 6: [9, 12] + - 6 6: [10, 0] + - 6 6: [10, 1] + - 6 6: [10, 2] + - 6 6: [10, 3] + - 6 6: [10, 4] + - 6 6: [10, 5] + - 6 6: [10, 6] + - 6 6: [10, 7] + - 6 6: [10, 8] + - 6 6: [10, 9] + - 6 6: [10, 10] + - 6 6: [10, 11] + - 6 6: [10, 12] + - 6 6: [11, 0] + - 6 6: [11, 1] + - 6 6: [11, 2] + - 6 6: [11, 3] + - 6 6: [11, 4] + - 6 6: [11, 5] + - 6 6: [11, 6] + - 6 6: [11, 7] + - 6 6: [11, 8] + - 6 6: [11, 9] + - 6 6: [11, 10] + - 6 6: [11, 11] + - 6 6: [11, 12] + - 6 6: [12, 0] + - 6 6: [12, 1] + - 6 6: [12, 2] + - 6 6: [12, 3] + - 6 6: [12, 4] + - 6 6: [12, 5] + - 6 6: [12, 6] + - 6 6: [12, 7] + - 6 6: [12, 8] + - 6 6: [12, 9] + - 6 6: [12, 10] + - 6 6: [12, 11] + - 6 6: [12, 12] + o_component: H + energy_component: summation + max_element: -1 + statistics: false + normalizer: false + boxcox: false + graph: + radius: 6.0 + max_num_nbr: 0 + create_from_dft: false + if_lcmp_graph: true + separate_onsite: false + new_sp: false + train: + epochs: 5 + pretrained: '' + resume: '' + train_ratio: 0.6 + val_ratio: 0.2 + test_ratio: 0.2 + early_stopping_loss: 0.0 + early_stopping_loss_epoch: + - 0.0 + - 500 + revert_then_decay: true + revert_threshold: 30 + revert_decay_epoch: + - 800 + - 2000 + - 3000 + - 4000 + revert_decay_gamma: + - 0.4 + - 0.5 + - 0.5 + - 0.4 + clip_grad: true + clip_grad_value: 4.2 + switch_sgd: false + switch_sgd_lr: 0.0001 + switch_sgd_epoch: -1 + hyperparameter: + batch_size: 4 + dtype: float32 + optimizer: adam + learning_rate: 0.001 + lr_scheduler: '' + lr_milestones: [] + momentum: 0.9 + weight_decay: 0 + criterion: MaskMSELoss + retain_edge_fea: true + lambda_eij: 0.0 + lambda_ei: 0.1 + lambda_etot: 0.0 + network: + atom_fea_len: 64 + edge_fea_len: 128 + gauss_stop: 6.0 + num_l: 5 + aggr: add + distance_expansion: GaussianBasis + if_exp: true + if_multiplelinear: false + if_edge_update: true + if_lcmp: true + normalization: LayerNorm + atom_update_net: CGConv + trainable_gaussians: false + type_affine: false +Dataset: + train: + dataset: + __class_name__: DeepHDataset + __init_params__: + split: train + split_seed: 42 + build_structure_cfg: + format: array + primitive: false + niggli: false + canocial: true + build_graph_cfg: + __class_name__: FindPointsInSpheres + __init_params__: + cutoff: 6.0 + pbc: + - 1 + - 1 + - 1 + eps: 1.0e-08 + config: ${DeepH} + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: true + drop_last: false + batch_size: 1 + loader: + num_workers: 0 + use_shared_memory: false + collate_fn: DefaultCollator + val: + dataset: + __class_name__: DeepHDataset + __init_params__: + split: val + split_seed: 42 + build_structure_cfg: + format: array + primitive: false + niggli: false + canocial: true + build_graph_cfg: + __class_name__: FindPointsInSpheres + __init_params__: + cutoff: 6.0 + pbc: + - 1 + - 1 + - 1 + eps: 1.0e-08 + config: ${DeepH} + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: false + drop_last: false + batch_size: 1 + loader: + num_workers: 0 + use_shared_memory: false + collate_fn: DefaultCollator + test: + dataset: + __class_name__: DeepHDataset + __init_params__: + split: test + split_seed: 42 + build_structure_cfg: + format: array + primitive: false + niggli: false + canocial: true + build_graph_cfg: + __class_name__: FindPointsInSpheres + __init_params__: + cutoff: 6.0 + pbc: + - 1 + - 1 + - 1 + eps: 1.0e-08 + config: ${DeepH} + sampler: + __class_name__: BatchSampler + __init_params__: + shuffle: false + drop_last: false + batch_size: 1 + loader: + num_workers: 0 + use_shared_memory: false + collate_fn: DefaultCollator diff --git a/electronic_structure/docs/deeph_fig2.png b/electronic_structure/docs/deeph_fig2.png new file mode 100644 index 00000000..3d41da9c Binary files /dev/null and b/electronic_structure/docs/deeph_fig2.png differ diff --git a/ppmat/calculator/deeph.py b/ppmat/calculator/deeph.py new file mode 100644 index 00000000..5ce80bf2 --- /dev/null +++ b/ppmat/calculator/deeph.py @@ -0,0 +1,385 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import glob +import importlib +import json +import os +import re +from pathlib import Path + +import numpy as np +from ase import Atoms +from ase.calculators.calculator import all_changes +from ase.data import atomic_numbers +from ase.io import write +from omegaconf import OmegaConf +from pymatgen.io.ase import AseAtomsAdaptor + +from ppmat.utils import logger + + +def _import_h5py(): + try: + import h5py + except ImportError as exc: + raise ImportError( + "OpenMX overlap processing requires h5py. Install the dependencies " + "from requirements.txt." + ) from exc + return h5py + + +def _parse_openmx_metadata(output_path): + lines = Path(output_path).read_text().splitlines() + orbital_by_element = {} + lattice = None + elements = [] + frac_coords = [] + + in_species = False + for index, line in enumerate(lines): + stripped = line.strip() + if stripped.startswith(""): + in_species = False + continue + if in_species: + fields = stripped.split() + if len(fields) >= 2: + basis = fields[1].rsplit("-", 1)[-1] + orbital_tokens = re.findall(r"([spdf])(\d+)", basis) + if orbital_tokens: + angular_momentum = {"s": 0, "p": 1, "d": 2, "f": 3} + orbital_by_element[fields[0]] = [ + angular_momentum[name] + for name, count in orbital_tokens + for _ in range(int(count)) + ] + + if "Atoms.UnitVectors.Unit" in line: + fields = stripped.split() + if len(fields) < 2 or fields[1].lower() != "ang": + raise ValueError("OpenMX lattice vectors must use Angstrom units.") + vectors = [] + cursor = index + 1 + while cursor < len(lines) and len(vectors) < 3: + values = lines[cursor].split() + if len(values) == 3: + try: + vectors.append([float(value) for value in values]) + except ValueError: + pass + cursor += 1 + if len(vectors) != 3: + raise ValueError("Can not parse lattice vectors from OpenMX output.") + lattice = np.asarray(vectors, dtype=np.float64) + + if "Fractional coordinates of the final structure" in line: + cursor = index + 1 + started = False + while cursor < len(lines): + fields = lines[cursor].split() + if len(fields) == 5 and fields[0].isdigit(): + expected_index = len(elements) + 1 + if int(fields[0]) != expected_index: + raise ValueError("Unexpected atom index in OpenMX output.") + elements.append(fields[1]) + frac_coords.append([float(value) for value in fields[2:5]]) + started = True + elif started: + break + cursor += 1 + + if lattice is None: + raise ValueError("Can not find lattice vectors in OpenMX output.") + if not elements: + raise ValueError("Can not find final fractional coordinates in OpenMX output.") + missing_orbitals = sorted(set(elements) - set(orbital_by_element)) + if missing_orbitals: + raise ValueError( + "Can not find OpenMX orbital definitions for: " f"{missing_orbitals}" + ) + return lattice, elements, np.asarray(frac_coords), orbital_by_element + + +def _read_openmx_overlaps(raw_dir, overlap_pattern): + h5py = _import_h5py() + overlap_paths = sorted(glob.glob(os.path.join(raw_dir, overlap_pattern))) + if not overlap_paths: + raise FileNotFoundError( + f"No OpenMX overlap files match {overlap_pattern!r} in {raw_dir}." + ) + + overlaps = {} + for overlap_path in overlap_paths: + with h5py.File(overlap_path, "r") as overlap_file: + for key, value in overlap_file.items(): + if key in overlaps: + raise ValueError(f"Duplicate OpenMX overlap block: {key}") + overlaps[key] = value[...] + return overlaps + + +def _build_local_coordinates(overlap_keys, cart_coords, lattice, tolerance=1e-8): + neighbours = {index: [] for index in range(cart_coords.shape[0])} + parsed_keys = [] + for key_str in overlap_keys: + key = json.loads(key_str) + if len(key) != 5: + raise ValueError(f"Invalid DeepH overlap key: {key_str}") + shift = np.asarray(key[:3], dtype=np.int64) + atom_i = int(key[3]) - 1 + atom_j = int(key[4]) - 1 + vector = cart_coords[atom_j] + shift @ lattice - cart_coords[atom_i] + distance = float(np.linalg.norm(vector)) + item = (distance, atom_j, shift, vector, key) + neighbours.setdefault(atom_i, []).append(item) + parsed_keys.append((key_str, item)) + + for atom_i, atom_neighbours in neighbours.items(): + if not atom_neighbours: + raise ValueError(f"Atom {atom_i} has no OpenMX overlap blocks.") + atom_neighbours.sort(key=lambda item: item[0]) + if atom_neighbours[0][0] > tolerance: + raise ValueError(f"Atom {atom_i} has no onsite overlap block.") + + rotations = {} + for key_str, (_, _, _, vector_ij, key) in parsed_keys: + atom_i = int(key[3]) - 1 + atom_neighbours = neighbours[atom_i] + if np.linalg.norm(vector_ij) > tolerance: + axis_source = vector_ij + else: + nonzero = [item[3] for item in atom_neighbours if item[0] > tolerance] + if not nonzero: + raise ValueError(f"Atom {atom_i} has no non-onsite overlap blocks.") + axis_source = nonzero[0] + + second_axis_source = None + for _, _, _, candidate, _ in atom_neighbours: + cross = np.cross(axis_source, candidate) + if np.linalg.norm(cross) > tolerance: + second_axis_source = candidate + break + if second_axis_source is None: + raise ValueError( + "There is no linearly independent bond for DeepH local " + f"coordinates around atom {atom_i}." + ) + + axis_1 = axis_source / np.linalg.norm(axis_source) + cross = np.cross(axis_source, second_axis_source) + axis_2 = cross / np.linalg.norm(cross) + axis_3 = np.cross(axis_1, axis_2) + rotations[key_str] = np.stack([axis_1, axis_2, axis_3], axis=-1) + return rotations + + +def preprocess_openmx_overlap( + raw_dir, + output_dir, + output_filename="openmx.out", + overlap_pattern="output/overlaps_*.h5", +): + """Convert overlap-only OpenMX output into DeepH inference inputs.""" + raw_dir = os.path.abspath(raw_dir) + output_dir = os.path.abspath(output_dir) + os.makedirs(output_dir, exist_ok=True) + + overlaps = _read_openmx_overlaps(raw_dir, overlap_pattern) + lattice, elements, frac_coords, orbital_by_element = _parse_openmx_metadata( + os.path.join(raw_dir, output_filename) + ) + cart_coords = frac_coords @ lattice + + np.savetxt(os.path.join(output_dir, "lat.dat"), lattice.T) + np.savetxt(os.path.join(output_dir, "rlat.dat"), np.linalg.inv(lattice) * 2 * np.pi) + np.savetxt(os.path.join(output_dir, "site_positions.dat"), cart_coords.T) + np.savetxt( + os.path.join(output_dir, "element.dat"), + np.asarray([atomic_numbers[element] for element in elements]), + fmt="%d", + ) + with open(os.path.join(output_dir, "orbital_types.dat"), "w") as file: + for element in elements: + file.write(" ".join(str(value) for value in orbital_by_element[element])) + file.write("\n") + + rotations = _build_local_coordinates(overlaps, cart_coords, lattice) + np.savez(os.path.join(output_dir, "overlaps.npz"), **overlaps) + np.savez(os.path.join(output_dir, "rc.npz"), **rotations) + return output_dir + + +class DeepHDFTInferenceTask: + def __init__( + self, + species, + num_l, + orbital=None, + orbital_config_path=None, + output_dir="deeph_dft_results", + dtype="float32", + run_dft=True, + processed_dirs=None, + ): + if orbital is not None and OmegaConf.is_config(orbital): + orbital = OmegaConf.to_container(orbital, resolve=True) + self.orbital_config_path = orbital_config_path + self.orbital = orbital + self.species = list(species) + self.num_l = num_l + self.output_dir = output_dir + self.dtype = dtype + self.run_dft = run_dft + self.processed_dirs = processed_dirs + + def __call__(self, interface_obj, structures): + logger.info("Run OpenMX overlap calculation and Paddle DeepH inference.") + interface_obj.run_deeph_inference( + structures=structures, + orbital_config_path=self.orbital_config_path, + orbital=self.orbital, + species=self.species, + num_l=self.num_l, + output_dir=self.output_dir, + dtype=self.dtype, + run_dft=self.run_dft, + processed_dirs=self.processed_dirs, + ) + logger.info("All DeepH DFT inference tasks finished successfully.") + + +class DeepHDFTCalculator: + """Run overlap-only DFT through ASE and infer Hamiltonians with DeepH.""" + + def __init__( + self, + predictor, + calculator_cls="ase.calculators.openmx.OpenMX", + calculator_kwargs=None, + command=None, + output_filename="openmx.log", + overlap_pattern="output/overlaps_*.h5", + **kwargs, + ): + del kwargs + self.predictor = predictor + self.calculator_cls = calculator_cls + self.calculator_kwargs = dict(calculator_kwargs or {}) + self.command = command + self.output_filename = output_filename + self.overlap_pattern = overlap_pattern + + @staticmethod + def _to_ase(structure): + if isinstance(structure, Atoms): + return structure.copy() + return AseAtomsAdaptor.get_atoms(structure) + + def _build_calculator(self, dft_dir): + module_name, class_name = self.calculator_cls.rsplit(".", 1) + calculator_class = getattr(importlib.import_module(module_name), class_name) + kwargs = dict(self.calculator_kwargs) + kwargs.setdefault("label", os.path.join(dft_dir, "openmx")) + if self.command is not None: + kwargs.setdefault("command", self.command) + return calculator_class(**kwargs) + + def _run_dft(self, structure, dft_dir): + os.makedirs(dft_dir, exist_ok=True) + atoms = self._to_ase(structure) + write(os.path.join(dft_dir, "structure.xyz"), atoms) + calculator = self._build_calculator(dft_dir) + if hasattr(calculator, "write_input") and hasattr(calculator, "run"): + calculator.atoms = atoms.copy() + calculator.write_input( + atoms=atoms, + properties=[], + system_changes=all_changes, + ) + calculator.run() + else: + atoms.calc = calculator + atoms.get_potential_energy() + + def run_deeph_inference( + self, + structures, + species, + num_l, + output_dir, + dtype, + run_dft, + processed_dirs, + orbital=None, + orbital_config_path=None, + ): + if not run_dft and processed_dirs is None: + raise ValueError("processed_dirs is required when run_dft=False.") + if processed_dirs is not None and len(processed_dirs) != len(structures): + raise ValueError("processed_dirs and structures must have the same length.") + + if orbital is None: + if orbital_config_path is not None: + orbital = self.predictor.load_orbital_config(orbital_config_path) + else: + orbital = self.predictor.config["DeepH"]["basic"]["orbital"] + os.makedirs(output_dir, exist_ok=True) + for index, structure in enumerate(structures): + sample_dir = os.path.join(output_dir, f"{index:06d}") + processed_dir = os.path.join(sample_dir, "processed") + if run_dft: + dft_dir = os.path.join(sample_dir, "dft") + self._run_dft(structure, dft_dir) + preprocess_openmx_overlap( + raw_dir=dft_dir, + output_dir=processed_dir, + output_filename=self.output_filename, + overlap_pattern=self.overlap_pattern, + ) + else: + processed_dir = os.path.abspath(processed_dirs[index]) + + required_files = { + "lat.dat", + "site_positions.dat", + "element.dat", + "orbital_types.dat", + "rc.npz", + } + missing = sorted( + filename + for filename in required_files + if not os.path.exists(os.path.join(processed_dir, filename)) + ) + if missing: + raise FileNotFoundError( + f"Incomplete DeepH DFT data in {processed_dir}: {missing}" + ) + + os.makedirs(sample_dir, exist_ok=True) + self.predictor.predict_hamiltonian( + processed_dir=processed_dir, + output_path=os.path.join(sample_dir, "rh_pred.npz"), + orbital=orbital, + species=species, + num_l=num_l, + dtype=dtype, + ) diff --git a/ppmat/datasets/__init__.py b/ppmat/datasets/__init__.py index c74d46d7..e2663d2a 100644 --- a/ppmat/datasets/__init__.py +++ b/ppmat/datasets/__init__.py @@ -42,7 +42,8 @@ from ppmat.datasets.msd_nmr_dataset import MSDnmrDataset from ppmat.datasets.msd_nmr_dataset import MSDnmrinfos from ppmat.datasets.density_dataset import DensityDataset -from ppmat.datasets.small_density_dataset import SmallDensityDataset +from ppmat.datasets.deeph_dataset import DeepHDataset +from ppmat.datasets.small_density_dataset import SmallDensityDataset from ppmat.datasets.sfin_dataset import SFINDataset from ppmat.datasets.num_atom_crystal_dataset import NumAtomsCrystalDataset from ppmat.datasets.oc20_s2ef_dataset import OC20S2EFDataset # noqa @@ -67,6 +68,7 @@ "MSDnmrDataset", "MatbenchDataset", "DensityDataset", + "DeepHDataset", "SmallDensityDataset", "SFINDataset", "OMol25Dataset", diff --git a/ppmat/datasets/deeph_dataset.py b/ppmat/datasets/deeph_dataset.py new file mode 100644 index 00000000..bd47b13c --- /dev/null +++ b/ppmat/datasets/deeph_dataset.py @@ -0,0 +1,527 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +import os.path as osp +import pickle +import time +from configparser import ConfigParser +from typing import Any +from typing import Dict +from typing import Tuple + +import numpy as np +import paddle +from omegaconf import OmegaConf + +from ppmat.datasets.build_structure import BuildStructure +from ppmat.datasets.custom_data_type import ConcatData +from ppmat.models import build_graph_converter +from ppmat.utils import logger +from ppmat.utils.crystal import lattices_to_params_shape_numpy +from ppmat.utils.misc import is_equal + + +def _as_tuple(config_files) -> Tuple[str, ...]: + if config_files is None: + return () + if isinstance(config_files, str): + return (config_files,) + return tuple(config_files) + + +def _load_config(config_files=None, config=None) -> ConfigParser: + if config is not None: + config = _to_container(config) + parser = ConfigParser() + for section, values in config.items(): + parser.add_section(section) + for key, value in values.items(): + if isinstance(value, (dict, list, tuple)): + value = json.dumps(value) + parser.set(section, key, str(value)) + return parser + + config = ConfigParser() + read_files = config.read(list(_as_tuple(config_files))) + if not read_files: + raise FileNotFoundError(f"Can not read DeepH config files: {config_files}") + return config + + +def _np_dtype(dtype: str): + if dtype == "float32": + return np.float32 + if dtype == "float16": + return np.float16 + if dtype == "float64": + return np.float64 + raise ValueError(f"Unknown DeepH dtype: {dtype}") + + +def _to_container(config): + if config is None: + return None + if OmegaConf.is_config(config): + return OmegaConf.to_container(config, resolve=True) + return config + + +def _list_structure_folders(raw_data_dir: str, interface: str, nums: int | None = None): + if interface != "npz": + raise NotImplementedError( + "DeepHDataset currently supports the validated npz Hamiltonian format." + ) + folder_list = sorted( + root for root, _, files in os.walk(raw_data_dir) if "rc.npz" in files + ) + if nums is not None: + folder_list = folder_list[:nums] + return folder_list + + +def _load_structure_arrays(folder: str): + lattice = np.loadtxt(os.path.join(folder, "lat.dat")).T + atom_types = np.loadtxt(os.path.join(folder, "element.dat")).astype(int).tolist() + cart_coords = np.loadtxt(os.path.join(folder, "site_positions.dat")).T + frac_coords = cart_coords @ np.linalg.inv(lattice) + lengths, angles = lattices_to_params_shape_numpy(lattice) + return { + "lengths": lengths, + "angles": angles, + "atom_types": atom_types, + "cart_coords": cart_coords, + "frac_coords": frac_coords, + } + + +class DeepHDataset(paddle.io.Dataset): + """DeepH dataset adapter for PaddleMaterials. + + This adapter exposes PaddleMaterials geometric `Data` objects and keeps the + LCMP subgraph metadata required by DeepH. + """ + + _CACHE: Dict[Tuple, Dict] = {} + _STRUCTURE_CACHE: Dict[Tuple[str, str], Dict] = {} + + def __init__( + self, + split: str, + config_files=None, + config=None, + nums: int | None = None, + split_seed: int | None = None, + build_structure_cfg: Dict[str, Any] | None = None, + build_graph_cfg: Dict[str, Any] | None = None, + cache_path: str | None = None, + overwrite: bool = False, + ): + super().__init__() + if split not in {"train", "val", "test"}: + raise ValueError(f"Unsupported split '{split}'. Expected train/val/test.") + if config is None and config_files is None: + raise ValueError("Either `config` or `config_files` must be provided.") + + self.config_files = _as_tuple(config_files) + self.config = _to_container(config) + self.split = split + self.nums = nums + self.split_seed = int(split_seed) if split_seed is not None else None + self.build_structure_cfg = _to_container(build_structure_cfg) or { + "format": "array", + "primitive": False, + "niggli": False, + "canocial": True, + } + self.build_graph_cfg = _to_container(build_graph_cfg) + self.cache_path = cache_path + self.overwrite = overwrite + + shared = self._load_shared() + self.dataset = shared["dataset"] + self.folder_list = shared["folder_list"] + self.indices = shared["split_indices"][split] + self.info = shared["info"] + + def _load_shared(self) -> Dict: + cache_key = ( + self.config_files, + json.dumps(self.config, sort_keys=True), + self.nums, + self.split_seed or -1, + json.dumps(self.build_structure_cfg, sort_keys=True), + json.dumps(self.build_graph_cfg, sort_keys=True), + self.cache_path or "", + int(self.overwrite), + ) + shared = self._CACHE.get(cache_key) + if shared is not None: + return shared + + config = _load_config(self.config_files, self.config) + folder_list = _list_structure_folders( + config.get("basic", "raw_dir"), + config.get("basic", "interface"), + self.nums, + ) + + dataset, dataset_info = self._load_or_build_graphs(config, folder_list) + + target = config.get("basic", "target") + if target != "hamiltonian": + raise NotImplementedError( + "DeepHDataset currently supports the validated hamiltonian target." + ) + orbital = json.loads(config.get("basic", "orbital")) + dataset = self._make_hamiltonian_mask( + dataset, + orbital=orbital, + num_orbital=len(orbital), + spinful=dataset_info["spinful"], + index_to_Z=dataset_info["index_to_Z"], + ) + + dataset_size = len(dataset) + train_size = int(config.getfloat("train", "train_ratio") * dataset_size) + val_size = int(config.getfloat("train", "val_ratio") * dataset_size) + test_size = int(config.getfloat("train", "test_ratio") * dataset_size) + + seed = ( + self.split_seed + if self.split_seed is not None + else config.getint("basic", "seed", fallback=42) + ) + rng = np.random.RandomState(seed) + indices = list(range(dataset_size)) + rng.shuffle(indices) + split_indices = { + "train": indices[:train_size], + "val": indices[train_size : train_size + val_size], + "test": indices[train_size + val_size : train_size + val_size + test_size], + } + + shared = { + "dataset": dataset, + "folder_list": folder_list, + "split_indices": split_indices, + "info": { + "num_species": len(dataset_info["index_to_Z"]), + "spinful": dataset_info["spinful"], + "dataset_size": dataset_size, + "config_files": list(self.config_files), + }, + } + self._CACHE[cache_key] = shared + return shared + + def _graph_cache_path(self, config) -> str: + dataset_name = config.get("basic", "dataset_name") + interface = config.get("basic", "interface") + num_l = config.getint("network", "num_l") + radius = config.getfloat("graph", "radius") + max_num_nbr = config.getint("graph", "max_num_nbr") + suffix = f"{radius}r{max_num_nbr}mn" + if config.getboolean("graph", "create_from_DFT", fallback=True): + suffix = "FromDFT" + nums_suffix = "all" if self.nums is None else f"n{self.nums}" + if self.cache_path is not None: + return self.cache_path + return osp.join( + config.get("basic", "graph_dir"), + f"PPMatDeepHGraph-{interface}-{dataset_name}-{num_l}l-{suffix}-{nums_suffix}", + ) + + def _load_or_build_graphs(self, config, folder_list): + cache_path = self._graph_cache_path(config) + graph_cache_path = osp.join(cache_path, "graphs.pkl") + cfg_cache_path = osp.join(cache_path, "dataset_cfg.pkl") + expected_cfg = self._cache_cfg(config, folder_list) + + if osp.exists(graph_cache_path) and not self.overwrite: + try: + cached_cfg = self.load_from_cache(cfg_cache_path) + if is_equal(cached_cfg, expected_cfg): + loaded = self.load_from_cache(graph_cache_path) + return loaded["graphs"], loaded["info"] + logger.warning("DeepH cache config differs from current settings.") + except Exception as e: + logger.warning(e) + logger.warning("Failed to load DeepH graph cache, will rebuild.") + + begin = time.time() + graphs = [ + self._build_graph_from_folder(folder, config) for folder in folder_list + ] + index_to_Z, Z_to_index = self._element_statistics(graphs) + spinful = bool(graphs[0].spinful) + for graph in graphs: + assert spinful == graph.spinful + + info = { + "spinful": spinful, + "index_to_Z": index_to_Z, + "Z_to_index": Z_to_index, + } + self.save_to_cache(cfg_cache_path, expected_cfg) + self.save_to_cache(graph_cache_path, {"graphs": graphs, "info": info}) + logger.info( + "Finish building PaddleMaterials graph cache with " + f"{len(graphs)} structures, " + f"cost {time.time() - begin:.0f} seconds" + ) + return graphs, info + + def _cache_cfg(self, config, folder_list): + return { + "config_files": list(self.config_files), + "folders": list(folder_list), + "nums": self.nums, + "build_structure_cfg": self.build_structure_cfg, + "build_graph_cfg": self._graph_converter_cfg(config), + "interface": config.get("basic", "interface"), + "target": config.get("basic", "target"), + "num_l": config.getint("network", "num_l"), + "dtype": config.get("hyperparameter", "dtype"), + "radius": config.getfloat("graph", "radius"), + "max_num_nbr": config.getint("graph", "max_num_nbr"), + } + + def _graph_converter_cfg(self, config): + if self.build_graph_cfg is not None: + return self.build_graph_cfg + return { + "__class_name__": "FindPointsInSpheres", + "__init_params__": { + "cutoff": config.getfloat("graph", "radius"), + "pbc": (1, 1, 1), + "eps": 1e-8, + }, + } + + def _deeph_graph_converter_cfg(self, config): + return { + "__class_name__": "DeepHGraphConverter", + "__init_params__": { + "radius": config.getfloat("graph", "radius"), + "max_num_nbr": config.getint("graph", "max_num_nbr"), + "default_dtype": _np_dtype(config.get("hyperparameter", "dtype")), + "interface": config.get("basic", "interface"), + "num_l": config.getint("network", "num_l"), + "create_from_DFT": config.getboolean( + "graph", "create_from_DFT", fallback=True + ), + "if_lcmp_graph": config.getboolean( + "graph", "if_lcmp_graph", fallback=True + ), + "separate_onsite": config.getboolean( + "graph", "separate_onsite", fallback=False + ), + "target": config.get("basic", "target"), + "huge_structure": False, + "if_new_sp": config.getboolean("graph", "new_sp", fallback=False), + }, + } + + def _build_graph_from_folder(self, folder: str, config): + structure_info = self._load_structure(folder) + structure = structure_info["structure"] + create_from_dft = config.getboolean("graph", "create_from_DFT", fallback=True) + converter_graph = None + if not create_from_dft: + structure_graph_converter = build_graph_converter( + self._graph_converter_cfg(config) + ) + converter_graph = structure_graph_converter(structure) + if converter_graph is None: + raise ValueError(f"Failed to build graph for DeepH structure: {folder}") + deeph_converter = build_graph_converter(self._deeph_graph_converter_cfg(config)) + return deeph_converter(structure, folder, converter_graph) + + @staticmethod + def _make_hamiltonian_mask(dataset, orbital, num_orbital, spinful, index_to_Z): + dataset_mask = [] + for data in dataset: + oij_value = data.term_real + if not np.all(data.term_mask): + raise NotImplementedError( + "Graph radius including hopping without calculation is not " + "supported." + ) + + if spinful: + out_fea_len = num_orbital * 8 + else: + out_fea_len = num_orbital + + mask = np.zeros((data.edge_attr.shape[0], out_fea_len), dtype=np.int8) + label = np.zeros( + (data.edge_attr.shape[0], out_fea_len), dtype=oij_value.dtype + ) + atomic_number_edge_i = index_to_Z[data.x[data.edge_index[0]]] + atomic_number_edge_j = index_to_Z[data.x[data.edge_index[1]]] + + for index_out, orbital_dict in enumerate(orbital): + for n_m_str, a_b in orbital_dict.items(): + condition_atomic_number_i, condition_atomic_number_j = map( + int, n_m_str.split() + ) + condition_orbital_i, condition_orbital_j = a_b + condition = (atomic_number_edge_i == condition_atomic_number_i) & ( + atomic_number_edge_j == condition_atomic_number_j + ) + + if spinful: + mask[:, 8 * index_out : 8 * (index_out + 1)] = np.where( + condition[:, None], + 1, + 0, + ) + else: + mask[:, index_out] += np.where(condition, 1, 0) + + if spinful: + value = oij_value[:, condition_orbital_i, condition_orbital_j] + label[:, 8 * index_out : 8 * (index_out + 1)] = np.where( + condition[:, None], + value, + np.zeros_like(value), + ) + else: + label[:, index_out] += np.where( + condition, + oij_value[:, condition_orbital_i, condition_orbital_j], + np.zeros(data.edge_attr.shape[0], dtype=oij_value.dtype), + ) + + assert len(np.where((mask != 1) & (mask != 0))[0]) == 0 + data.mask = mask.astype(bool) + del data.term_mask + data.label = label + del data.term_real + dataset_mask.append(data) + return dataset_mask + + @staticmethod + def _element_statistics(data_list): + index_to_Z = np.unique(data_list[0].x).astype(np.int64) + Z_to_index = np.full((100,), -1, dtype=np.int64) + Z_to_index[index_to_Z] = np.arange(len(index_to_Z), dtype=np.int64) + for data in data_list: + data.x = Z_to_index[data.x] + return index_to_Z, Z_to_index + + def _load_structure(self, folder: str) -> Dict: + structure_cache_key = ( + folder, + json.dumps(self.build_structure_cfg, sort_keys=True), + ) + cached = self._STRUCTURE_CACHE.get(structure_cache_key) + if cached is not None: + return cached + + arrays = _load_structure_arrays(folder) + builder = BuildStructure(**self.build_structure_cfg) + structure = BuildStructure.build_one( + { + "lengths": arrays["lengths"], + "angles": arrays["angles"], + "frac_coords": arrays["frac_coords"], + "atom_types": arrays["atom_types"], + }, + builder.format, + builder.primitive, + builder.niggli, + builder.canocial, + ) + cached = { + "folder": folder, + "structure": structure, + "lattice": np.asarray(structure.lattice.matrix, dtype=np.float32), + "frac_coords": np.asarray(structure.frac_coords, dtype=np.float32), + "cart_coords": np.asarray(structure.cart_coords, dtype=np.float32), + "atomic_numbers": np.asarray(structure.atomic_numbers, dtype=np.int64), + } + self._STRUCTURE_CACHE[structure_cache_key] = cached + return cached + + @staticmethod + def save_to_cache(cache_path: str, data: Any): + os.makedirs(osp.dirname(cache_path), exist_ok=True) + with open(cache_path, "wb") as f: + pickle.dump(data, f) + + @staticmethod + def load_from_cache(cache_path: str): + if not osp.exists(cache_path): + raise FileNotFoundError(f"No such file or directory: {cache_path}") + with open(cache_path, "rb") as f: + return pickle.load(f) + + def __len__(self): + return len(self.indices) + + def __getitem__(self, index: int) -> Dict: + graph_index = self.indices[index] + graph = self.dataset[graph_index] + folder = self.folder_list[graph_index] + structure_info = self._load_structure(folder) + + if hasattr(graph, "subgraph_dict") and graph.subgraph_dict is not None: + subgraph_dict = graph.subgraph_dict + elif hasattr(graph, "subgraph") and graph.subgraph is not None: + subgraph = graph.subgraph + subgraph_dict = { + "subgraph_atom_idx": subgraph[0], + "subgraph_edge_idx": subgraph[1], + "subgraph_edge_ang": subgraph[2], + "subgraph_index": subgraph[3], + } + else: + raise ValueError("DeepH sample does not contain LCMP subgraph metadata.") + + return { + "x": ConcatData(np.asarray(graph.x, dtype=np.int64)), + "edge_index": ConcatData(np.asarray(graph.edge_index, dtype=np.int64)), + "edge_attr": ConcatData(np.asarray(graph.edge_attr, dtype=np.float32)), + "batch": ConcatData(np.zeros(graph.x.shape[0], dtype=np.int64)), + "label": ConcatData(np.asarray(graph.label, dtype=np.float32)), + "mask": ConcatData(np.asarray(graph.mask, dtype=bool)), + "pos": ConcatData( + np.asarray(structure_info["cart_coords"], dtype=np.float32) + ), + "sub_atom_idx": ConcatData( + np.asarray(subgraph_dict["subgraph_atom_idx"], dtype=np.int64) + ), + "sub_edge_idx": ConcatData( + np.asarray(subgraph_dict["subgraph_edge_idx"], dtype=np.int64) + ), + "sub_edge_ang": ConcatData( + np.asarray(subgraph_dict["subgraph_edge_ang"], dtype=np.float32) + ), + "sub_index": ConcatData( + np.asarray(subgraph_dict["subgraph_index"], dtype=np.int64) + ), + "structure_lattice": ConcatData( + np.asarray(structure_info["lattice"], dtype=np.float32).reshape(1, 3, 3) + ), + "structure_frac_coords": ConcatData( + np.asarray(structure_info["frac_coords"], dtype=np.float32) + ), + "structure_atomic_numbers": ConcatData( + np.asarray(structure_info["atomic_numbers"], dtype=np.int64) + ), + "structure_folder": structure_info["folder"], + } diff --git a/ppmat/datasets/deeph_graph.py b/ppmat/datasets/deeph_graph.py new file mode 100644 index 00000000..27666353 --- /dev/null +++ b/ppmat/datasets/deeph_graph.py @@ -0,0 +1,561 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DeepH Hamiltonian graph construction for PaddleMaterials. + +This module implements the DeepH graphene training and inference graph paths. +Training can use radius graphs or the sparsity pattern of DFT matrices. Inference +uses the overlap-derived local-coordinate keys produced by the DFT workflow. +""" + +import itertools +import json +import os + +import numpy as np + + +class Data: + """Small array container used before conversion to PaddleMaterials Data.""" + + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + +def _semifactorial(x): + y = 1.0 + for n in range(x, 1, -2): + y *= n + return y + + +def _pochhammer(x, k): + xf = float(x) + for n in range(x + 1, x + k): + xf *= n + return xf + + +class _SphericalHarmonics: + """Real spherical harmonics used by the LCMP subgraph angular features.""" + + def __init__(self): + self.leg = {} + + def clear(self): + self.leg = {} + + def _negative_lpmv(self, degree, order, value): + if order < 0: + value *= (-1) ** order / _pochhammer(degree + order + 1, -2 * order) + return value + + def lpmv(self, degree, order, x): + order_abs = abs(order) + if (degree, order) in self.leg: + return self.leg[(degree, order)] + if order_abs > degree: + return None + if degree == 0: + self.leg[(degree, order)] = np.ones_like(x) + return self.leg[(degree, order)] + if order_abs == degree: + value = (-1) ** order_abs * _semifactorial(2 * order_abs - 1) + value *= np.power(1 - x * x, order_abs / 2) + self.leg[(degree, order)] = self._negative_lpmv(degree, order, value) + return self.leg[(degree, order)] + + self.lpmv(degree - 1, order, x) + value = ( + ((2 * degree - 1) / (degree - order_abs)) + * x + * self.lpmv(degree - 1, order_abs, x) + ) + if degree - order_abs > 1: + value -= ((degree + order_abs - 1) / (degree - order_abs)) * self.leg[ + (degree - 2, order_abs) + ] + if order < 0: + value = self._negative_lpmv(degree, order, value) + self.leg[(degree, order)] = value + return self.leg[(degree, order)] + + def get_element(self, degree, order, theta, phi): + norm = np.sqrt((2 * degree + 1) / (4 * np.pi)) + leg = self.lpmv(degree, abs(order), np.cos(theta)) + if order == 0: + return norm * leg + if order > 0: + value = np.cos(order * phi) * leg + else: + value = np.sin(abs(order) * phi) * leg + norm *= np.sqrt( + 2.0 + / _pochhammer( + degree - abs(order) + 1, + 2 * abs(order), + ) + ) + return value * norm + + def get(self, degree, theta, phi, refresh=True): + if refresh: + self.clear() + return np.stack( + [ + self.get_element(degree, order, theta, phi) + for order in range(-degree, degree + 1) + ], + axis=-1, + ) + + +def _get_spherical_from_cartesian(cartesian): + spherical = np.zeros(cartesian.shape[:-1] + (2,), dtype=cartesian.dtype) + r_xy = cartesian[..., 1] ** 2 + cartesian[..., 2] ** 2 + spherical[..., 0] = np.arctan2(np.sqrt(r_xy), cartesian[..., 0]) + spherical[..., 1] = np.arctan2(cartesian[..., 2], cartesian[..., 1]) + return spherical + + +def _load_orbital_types(path): + orbital_types = [] + with open(path) as f: + line = f.readline() + while line: + orbital_types.append(list(map(int, line.split()))) + line = f.readline() + return [ + sum(2 * orbital + 1 for orbital in atom_orbitals) + for atom_orbitals in orbital_types + ] + + +def _validate_supported_path( + interface, + target, + create_from_DFT, + if_lcmp_graph, + separate_onsite, + if_new_sp, + huge_structure, +): + if interface not in {"npz", "npz_rc_only"}: + raise NotImplementedError( + "DeepH PaddleMaterials integration supports interface='npz' and " + "'npz_rc_only'." + ) + if target != "hamiltonian": + raise NotImplementedError( + "DeepH PaddleMaterials integration currently supports target='hamiltonian'." + ) + if interface == "npz_rc_only" and not create_from_DFT: + raise ValueError("interface='npz_rc_only' requires create_from_DFT=True.") + if not if_lcmp_graph: + raise NotImplementedError( + "DeepH PaddleMaterials integration expects if_lcmp_graph=True." + ) + if separate_onsite: + raise NotImplementedError( + "DeepH separate_onsite=True is not enabled in this integration." + ) + if if_new_sp: + raise NotImplementedError( + "DeepH new_sp=True is not enabled in this integration." + ) + if huge_structure: + raise NotImplementedError( + "DeepH huge_structure=True is not enabled in this integration." + ) + + +def _periodic_image_ranges(frac_coords, lattice, radius): + reciprocal_lattice = np.linalg.inv(lattice).T * 2 * np.pi + recp_len = np.sqrt(np.sum(reciprocal_lattice**2, axis=1)) + maxr = np.ceil((radius + 0.15) * recp_len / (2 * np.pi)) + nmin = np.floor(np.min(frac_coords, axis=0)) - maxr + nmax = np.ceil(np.max(frac_coords, axis=0)) + maxr + return [np.arange(x, y, dtype="int64") for x, y in zip(nmin, nmax)] + + +def _build_edges_from_converter_graph(converter_graph, default_dtype): + edge_idx = np.asarray(converter_graph.edges, dtype=np.int64).T + cart_coords = np.asarray( + converter_graph.node_feat["cart_coords"], + dtype=default_dtype, + ) + lattice = np.asarray( + converter_graph.node_feat["lattice"], + dtype=default_dtype, + ).reshape(3, 3) + pbc_offset = np.asarray( + converter_graph.edge_feat["pbc_offset"], + dtype=default_dtype, + ) + edge_dist = np.asarray( + converter_graph.edge_feat["bond_dist"], + dtype=default_dtype, + ).reshape(-1, 1) + src, dst = edge_idx + dst_cart_periodic = cart_coords[dst] + pbc_offset @ lattice + edge_fea = np.concatenate( + [ + edge_dist, + cart_coords[src], + dst_cart_periodic, + cart_coords[dst], + ], + axis=-1, + ).astype(default_dtype) + onsite_src = np.arange(cart_coords.shape[0], dtype=np.int64) + onsite_edge_idx = np.stack([onsite_src, onsite_src]) + onsite_edge_fea = np.concatenate( + [ + np.zeros((cart_coords.shape[0], 1), dtype=default_dtype), + cart_coords, + cart_coords, + cart_coords, + ], + axis=-1, + ) + edge_idx = np.concatenate([edge_idx, onsite_edge_idx], axis=1) + edge_fea = np.concatenate([edge_fea, onsite_edge_fea], axis=0) + src, dst = edge_idx + + num_atom = cart_coords.shape[0] + atom_idx_connect, edge_idx_connect = [], [] + for atom_idx in range(num_atom): + outgoing_edge_idx = np.where(src == atom_idx)[0] + edge_idx_connect.append(outgoing_edge_idx) + atom_idx_connect.append(dst[outgoing_edge_idx]) + return edge_idx, edge_fea, atom_idx_connect, edge_idx_connect + + +def _load_npz_terms(tb_folder, default_dtype): + atom_num_orbital = _load_orbital_types(os.path.join(tb_folder, "orbital_types.dat")) + + read_terms = {} + hopping_dict_read = np.load(os.path.join(tb_folder, "rh.npz")) + for k, v in hopping_dict_read.items(): + key = json.loads(k) + key = (key[0], key[1], key[2], key[3] - 1, key[4] - 1) + read_terms[key] = np.asarray(v, dtype=default_dtype) + + local_rotation_dict = {} + local_rotation_dict_read = np.load(os.path.join(tb_folder, "rc.npz")) + for k, v in local_rotation_dict_read.items(): + key = json.loads(k) + key = (key[0], key[1], key[2], key[3] - 1, key[4] - 1) + local_rotation_dict[key] = np.asarray(v, dtype=default_dtype) + + return atom_num_orbital, read_terms, local_rotation_dict + + +def _load_npz_rotations(tb_folder, default_dtype): + atom_num_orbital = _load_orbital_types(os.path.join(tb_folder, "orbital_types.dat")) + rotations = {} + with np.load(os.path.join(tb_folder, "rc.npz")) as rotation_file: + for key_str, value in rotation_file.items(): + key = json.loads(key_str) + key = (key[0], key[1], key[2], key[3] - 1, key[4] - 1) + rotations[key] = np.asarray(value, dtype=default_dtype) + return atom_num_orbital, rotations + + +def _build_edges_from_rotation_keys( + cart_coords, + lattice, + local_rotation_dict, + default_dtype, +): + edge_keys = list(local_rotation_dict) + if not edge_keys: + raise ValueError("DeepH local-coordinate file does not contain any edges.") + + edge_idx = np.asarray([[key[3], key[4]] for key in edge_keys], dtype=np.int64).T + lattice_shifts = np.asarray([key[:3] for key in edge_keys], dtype=default_dtype) + src, dst = edge_idx + dst_cart_periodic = cart_coords[dst] + lattice_shifts @ lattice + edge_dist = np.linalg.norm( + dst_cart_periodic - cart_coords[src], axis=1, keepdims=True + ) + edge_fea = np.concatenate( + [ + edge_dist, + cart_coords[src], + dst_cart_periodic, + cart_coords[dst], + ], + axis=-1, + ).astype(default_dtype) + + atom_idx_connect = [] + edge_idx_connect = [] + for atom_idx in range(cart_coords.shape[0]): + outgoing_edge_idx = np.where(src == atom_idx)[0] + if outgoing_edge_idx.size == 0: + raise ValueError(f"Atom {atom_idx} has no overlap-derived DeepH edges.") + edge_idx_connect.append(outgoing_edge_idx) + atom_idx_connect.append(dst[outgoing_edge_idx]) + + local_rotation = np.stack( + [local_rotation_dict[key] for key in edge_keys], axis=0 + ).astype(default_dtype) + return ( + edge_idx, + edge_fea, + atom_idx_connect, + edge_idx_connect, + local_rotation, + ) + + +def _attach_terms( + edge_idx, + edge_fea, + lattice, + atom_num_orbital, + read_terms, + local_rotation_dict, + default_dtype, +): + max_num_orbital = max(atom_num_orbital) + term_mask = np.zeros(edge_fea.shape[0], dtype=bool) + term_real = np.full( + [edge_fea.shape[0], max_num_orbital, max_num_orbital], + np.nan, + dtype=default_dtype, + ) + local_rotation = [] + inv_lattice = np.linalg.inv(lattice).astype(default_dtype) + + for index_edge in range(edge_fea.shape[0]): + lattice_shift = ( + np.rint( + edge_fea[index_edge, 4:7] @ inv_lattice + - edge_fea[index_edge, 7:10] @ inv_lattice + ) + .astype(int) + .tolist() + ) + i, j = edge_idx[:, index_edge] + key_term = (*lattice_shift, int(i), int(j)) + if key_term not in read_terms: + raise NotImplementedError( + "Graph radius including hopping without calculation is not supported." + ) + term_mask[index_edge] = True + term_real[ + index_edge, + : atom_num_orbital[i], + : atom_num_orbital[j], + ] = read_terms[key_term] + local_rotation.append(local_rotation_dict[key_term]) + + return term_mask, term_real, np.stack(local_rotation, axis=0) + + +def _build_lcmp_subgraph( + edge_idx, edge_fea, local_rotation, atom_idx_connect, edge_idx_connect, num_l +): + r_vec = edge_fea[:, 1:4] - edge_fea[:, 4:7] + r_vec = np.matmul( + r_vec[:, None, None, :], + local_rotation[None, :, :, :], + ).reshape(-1, 3) + + r_vec_sp = _get_spherical_from_cartesian(r_vec) + sph_harm_func = _SphericalHarmonics() + angular_expansion = [] + for l_value in range(num_l): + angular_expansion.append( + sph_harm_func.get(l_value, r_vec_sp[:, 0], r_vec_sp[:, 1]) + ) + angular_expansion = np.concatenate(angular_expansion, axis=-1).reshape( + edge_fea.shape[0], + edge_fea.shape[0], + -1, + ) + + subgraph_atom_idx_list = [] + subgraph_edge_idx_list = [] + subgraph_edge_ang_list = [] + subgraph_index = [] + index_cursor = 0 + + for index in range(edge_fea.shape[0]): + i, j = edge_idx[:, index] + + subgraph_edge_idx = np.asarray(list(edge_idx_connect[i]), dtype=np.int64) + subgraph_atom_idx = np.stack( + [np.repeat(i, len(atom_idx_connect[i])), atom_idx_connect[i]], + axis=1, + ) + subgraph_atom_idx_list.append(subgraph_atom_idx) + subgraph_edge_idx_list.append(subgraph_edge_idx) + subgraph_edge_ang_list.append(angular_expansion[subgraph_edge_idx, index, :]) + subgraph_index += [index_cursor] * len(atom_idx_connect[i]) + index_cursor += 1 + + subgraph_edge_idx = np.asarray(list(edge_idx_connect[j]), dtype=np.int64) + subgraph_atom_idx = np.stack( + [np.repeat(j, len(atom_idx_connect[j])), atom_idx_connect[j]], + axis=1, + ) + subgraph_atom_idx_list.append(subgraph_atom_idx) + subgraph_edge_idx_list.append(subgraph_edge_idx) + subgraph_edge_ang_list.append(angular_expansion[subgraph_edge_idx, index, :]) + subgraph_index += [index_cursor] * len(atom_idx_connect[j]) + index_cursor += 1 + + return { + "subgraph_atom_idx": np.concatenate(subgraph_atom_idx_list, axis=0).astype( + np.int64 + ), + "subgraph_edge_idx": np.concatenate(subgraph_edge_idx_list, axis=0).astype( + np.int64 + ), + "subgraph_edge_ang": np.concatenate(subgraph_edge_ang_list, axis=0).astype( + edge_fea.dtype + ), + "subgraph_index": np.asarray(subgraph_index, dtype=np.int64), + } + + +def get_graph( + cart_coords, + frac_coords, + numbers, + stru_id, + r, + max_num_nbr, + numerical_tol, + lattice, + default_dtype, + tb_folder, + interface, + num_l, + create_from_DFT, + if_lcmp_graph, + separate_onsite, + target="hamiltonian", + huge_structure=False, + only_get_R_list=False, + if_new_sp=False, + **kwargs, +): + _validate_supported_path( + interface, + target, + create_from_DFT, + if_lcmp_graph, + separate_onsite, + if_new_sp, + huge_structure, + ) + if tb_folder is None: + raise ValueError("DeepH graph construction requires tb_folder.") + converter_graph = kwargs.get("converter_graph") + if max_num_nbr > 0: + raise NotImplementedError( + "DeepH PaddleMaterials integration currently relies on radius-based " + "FindPointsInSpheres graph conversion and expects max_num_nbr=0." + ) + + cart_coords = np.asarray(cart_coords, dtype=default_dtype) + frac_coords = np.asarray(frac_coords, dtype=default_dtype) + numbers = np.asarray(numbers, dtype=np.int64) + lattice = np.asarray(lattice, dtype=default_dtype) + + if only_get_R_list: + return np.array( + list(itertools.product(*_periodic_image_ranges(frac_coords, lattice, r))), + dtype=default_dtype, + ) + + term_mask = term_real = None + if create_from_DFT: + atom_num_orbital, local_rotation_dict = _load_npz_rotations( + tb_folder, default_dtype + ) + ( + edge_idx, + edge_fea, + atom_idx_connect, + edge_idx_connect, + local_rotation, + ) = _build_edges_from_rotation_keys( + cart_coords, + lattice, + local_rotation_dict, + default_dtype, + ) + if interface == "npz": + _, read_terms, _ = _load_npz_terms(tb_folder, default_dtype) + term_mask, term_real, _ = _attach_terms( + edge_idx, + edge_fea, + lattice, + atom_num_orbital, + read_terms, + local_rotation_dict, + default_dtype, + ) + else: + if converter_graph is None: + raise ValueError( + "Radius-based DeepH graph construction expects a " + "PaddleMaterials graph_converter graph." + ) + ( + edge_idx, + edge_fea, + atom_idx_connect, + edge_idx_connect, + ) = _build_edges_from_converter_graph(converter_graph, default_dtype) + atom_num_orbital, read_terms, local_rotation_dict = _load_npz_terms( + tb_folder, default_dtype + ) + term_mask, term_real, local_rotation = _attach_terms( + edge_idx, + edge_fea, + lattice, + atom_num_orbital, + read_terms, + local_rotation_dict, + default_dtype, + ) + subgraph = _build_lcmp_subgraph( + edge_idx, + edge_fea, + local_rotation, + atom_idx_connect, + edge_idx_connect, + num_l, + ) + + data = Data( + x=numbers, + edge_index=edge_idx, + edge_attr=edge_fea, + stru_id=stru_id, + onsite_term_real=None, + atom_num_orbital=np.asarray(atom_num_orbital, dtype=np.int64), + subgraph_dict=subgraph, + spinful=False, + ) + if term_mask is not None: + data.term_mask = term_mask + data.term_real = term_real + return data diff --git a/ppmat/losses/__init__.py b/ppmat/losses/__init__.py index e5258105..c0a1202c 100644 --- a/ppmat/losses/__init__.py +++ b/ppmat/losses/__init__.py @@ -19,10 +19,12 @@ from ppmat.losses.l1_loss import MAELoss from ppmat.losses.l1_loss import SmoothL1Loss from ppmat.losses.loss_warper import LossWarper +from ppmat.losses.mse_loss import MaskMSELoss from ppmat.losses.mse_loss import MSELoss __all__ = [ "MSELoss", + "MaskMSELoss", "L1Loss", "SmoothL1Loss", "MAELoss", diff --git a/ppmat/losses/mse_loss.py b/ppmat/losses/mse_loss.py index 92df488a..13f12673 100644 --- a/ppmat/losses/mse_loss.py +++ b/ppmat/losses/mse_loss.py @@ -38,3 +38,13 @@ def forward(self, pred, label) -> paddle.Tensor: loss = F.mse_loss(pred, label, self.reduction) return loss + + +class MaskMSELoss(nn.Layer): + r"""Mean squared error over valid entries selected by a boolean mask.""" + + def forward(self, pred, label, mask) -> paddle.Tensor: + diff = paddle.square(pred - label) + mask = paddle.cast(mask, pred.dtype) + denom = paddle.clip(paddle.sum(mask), min=1.0) + return paddle.sum(diff * mask) / denom diff --git a/ppmat/models/__init__.py b/ppmat/models/__init__.py index 32f1c729..a9eb1241 100644 --- a/ppmat/models/__init__.py +++ b/ppmat/models/__init__.py @@ -36,6 +36,8 @@ from ppmat.models.diffnmr.diffnmr import MolecularGraphFormer from ppmat.models.diffnmr.diffnmr import NMRNetCLIP from ppmat.models.dimenetpp.dimenetpp import DimeNetPlusPlus +from ppmat.models.deeph.deeph import DeepHHamiltonian +from ppmat.models.deeph.deeph_graph_converter import DeepHGraphConverter from ppmat.models.mattergen.mattergen import MatterGen from ppmat.models.mattergen.mattergen import MatterGenWithCondition from ppmat.models.mattersim.m3gnet import M3GNet @@ -69,6 +71,8 @@ "NMRNetCLIP", "DiffPrior", "DiffNMR", + "DeepHHamiltonian", + "DeepHGraphConverter", "InfGCN", "MatENO", "SFIN", diff --git a/ppmat/models/deeph/__init__.py b/ppmat/models/deeph/__init__.py new file mode 100644 index 00000000..05cecdb4 --- /dev/null +++ b/ppmat/models/deeph/__init__.py @@ -0,0 +1,7 @@ +from .deeph import DeepHHamiltonian +from .deeph_graph_converter import DeepHGraphConverter + +__all__ = [ + "DeepHHamiltonian", + "DeepHGraphConverter", +] diff --git a/ppmat/models/deeph/deeph.py b/ppmat/models/deeph/deeph.py new file mode 100644 index 00000000..b4fdd9c3 --- /dev/null +++ b/ppmat/models/deeph/deeph.py @@ -0,0 +1,286 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import paddle +from paddle import nn +import paddle.nn.functional as F + +from ppmat.losses import MaskMSELoss +from ppmat.utils.scatter import scatter_sum + + +class GraphLayerNorm(nn.Layer): + def __init__(self, in_channels, eps=1e-5): + super().__init__() + self.eps = eps + self.weight = self.create_parameter([in_channels], default_initializer=nn.initializer.Constant(1.0)) + self.bias = self.create_parameter([in_channels], default_initializer=nn.initializer.Constant(0.0)) + + def forward(self, x, batch): + batch = paddle.cast(batch, "int64") + batch_size = int(paddle.max(batch).item()) + 1 + num_channels = x.shape[-1] + + ones = paddle.ones([x.shape[0], 1], dtype=x.dtype) + count = scatter_sum(ones, batch, dim=0, dim_size=batch_size) + norm = count * float(num_channels) + + summed = scatter_sum(x, batch, dim=0, dim_size=batch_size) + mean = paddle.sum(summed, axis=-1, keepdim=True) / norm + x_centered = x - paddle.gather(mean, batch, axis=0) + + var = scatter_sum(x_centered * x_centered, batch, dim=0, dim_size=batch_size) + var = paddle.sum(var, axis=-1, keepdim=True) / norm + out = x_centered / paddle.sqrt(paddle.gather(var, batch, axis=0) + self.eps) + return out * self.weight + self.bias + + +def gaussian_smearing(distances, offset, widths, centered=False): + if not centered: + coeff = -0.5 / paddle.pow(widths, 2) + diff = distances.unsqueeze(-1) - offset + else: + coeff = -0.5 / paddle.pow(offset, 2) + diff = distances.unsqueeze(-1) + return paddle.exp(coeff * paddle.pow(diff, 2)) + + +class GaussianBasis(nn.Layer): + def __init__(self, start=0.0, stop=5.0, n_gaussians=50): + super().__init__() + offset = paddle.linspace(start, stop, n_gaussians) + widths = paddle.full(offset.shape, offset[1] - offset[0], dtype="float32") + self.register_buffer("offsets", offset) + self.register_buffer("width", widths) + + def forward(self, distances): + return gaussian_smearing(distances, self.offsets, self.width, centered=False) + + +class MLPBlock(nn.Layer): + def __init__(self, in_dim, hidden_dim, out_dim, final_activation=True): + super().__init__() + self.fc1 = nn.Linear(in_dim, hidden_dim) + self.fc2 = nn.Linear(hidden_dim, out_dim) + self.final_activation = final_activation + + def forward(self, x): + x = F.silu(self.fc1(x)) + x = self.fc2(x) + if self.final_activation: + x = F.silu(x) + return x + + +class CGConv(nn.Layer): + def __init__(self, channels, dim, normalization="LayerNorm", if_exp=False): + super().__init__() + self.if_exp = if_exp + self.lin_f = nn.Linear(channels * 2 + dim, channels) + self.lin_s = nn.Linear(channels * 2 + dim, channels) + self.ln = GraphLayerNorm(channels) if normalization == "LayerNorm" else None + + def forward(self, x, edge_index, edge_attr, batch, distance): + row = paddle.cast(edge_index[0], "int64") + col = paddle.cast(edge_index[1], "int64") + x_j = paddle.gather(x, row, axis=0) + x_i = paddle.gather(x, col, axis=0) + z = paddle.concat([x_i, x_j, edge_attr], axis=-1) + out = F.sigmoid(self.lin_f(z)) * F.softplus(self.lin_s(z)) + if self.if_exp: + sigma = 3.0 + n = 2.0 + out = out * paddle.exp(-(distance ** n) / (sigma ** n) / 2.0).reshape([-1, 1]) + out = scatter_sum(out, col, dim=0, dim_size=x.shape[0]) + if self.ln is not None: + out = self.ln(out, batch) + return out + x + + +class MPLayer(nn.Layer): + def __init__(self, atom_dim, edge_dim, out_edge_dim, if_exp, if_edge_update=True, normalization="LayerNorm"): + super().__init__() + self.cgconv = CGConv(atom_dim, edge_dim, normalization=normalization, if_exp=if_exp) + self.if_edge_update = if_edge_update + if if_edge_update: + final_activation = out_edge_dim != 169 + self.e_lin = MLPBlock(edge_dim + atom_dim * 2, 128, out_edge_dim, final_activation) + + def forward(self, atom_fea, edge_idx, edge_fea, batch, distance): + atom_fea = self.cgconv(atom_fea, edge_idx, edge_fea, batch, distance) + if self.if_edge_update: + row = paddle.cast(edge_idx[0], "int64") + col = paddle.cast(edge_idx[1], "int64") + edge_fea = self.e_lin( + paddle.concat( + [paddle.gather(atom_fea, row, axis=0), paddle.gather(atom_fea, col, axis=0), edge_fea], + axis=-1, + ) + ) + return atom_fea, edge_fea + return atom_fea + + +class LCMPLayer(nn.Layer): + def __init__(self, atom_dim, edge_dim, out_dim, num_l, if_exp=False): + super().__init__() + self.if_exp = if_exp + self.lin_f = nn.Linear(atom_dim * 2 + edge_dim, atom_dim) + self.lin_s = nn.Linear(atom_dim * 2 + edge_dim, atom_dim) + self.e_lin = MLPBlock(edge_dim + atom_dim * 2 - num_l ** 2, 128, out_dim, final_activation=False) + + def forward(self, atom_fea, edge_fea, sub_atom_idx, sub_edge_idx, sub_edge_ang, sub_index, distance): + num_edge = edge_fea.shape[0] + sub_atom_idx = paddle.cast(sub_atom_idx, "int64") + sub_edge_idx = paddle.cast(sub_edge_idx, "int64") + sub_index = paddle.cast(sub_index, "int64") + gathered_atoms = paddle.gather(atom_fea, sub_atom_idx.reshape([-1]), axis=0).reshape( + [sub_atom_idx.shape[0], 2, atom_fea.shape[1]] + ) + gathered_edge = paddle.gather(edge_fea, sub_edge_idx, axis=0) + z = paddle.concat([gathered_atoms[:, 0, :], gathered_atoms[:, 1, :], gathered_edge, sub_edge_ang], axis=-1) + out = F.sigmoid(self.lin_f(z)) * F.softplus(self.lin_s(z)) + if self.if_exp: + sigma = 3.0 + n = 2.0 + edge_distance = paddle.gather(distance, sub_edge_idx, axis=0) + out = out * paddle.exp(-(edge_distance ** n) / (sigma ** n) / 2.0).reshape([-1, 1]) + out = scatter_sum(out, sub_index, dim=0, dim_size=num_edge * 2) + out = paddle.reshape(out, [num_edge, 2, -1]) + out = self.e_lin(paddle.concat([out[:, 0, :], out[:, 1, :], edge_fea], axis=-1)) + return out + + +class DeepHHamiltonian(nn.Layer): + def __init__( + self, + num_species, + in_atom_fea_len, + in_edge_fea_len, + num_orbital, + num_l=5, + gauss_stop=6.0, + if_exp=True, + normalization="LayerNorm", + target_name="label", + loss_eps=1e-8, + **kwargs, + ): + super().__init__() + self.num_species = num_species + self.num_orbital = num_orbital + self.num_l = num_l + self.target_name = target_name + self.loss_eps = loss_eps + + self.embed = nn.Embedding(num_species + 5, in_atom_fea_len) + self.distance_expansion = GaussianBasis(0.0, gauss_stop, in_edge_fea_len) + self.mp1 = MPLayer(in_atom_fea_len, in_edge_fea_len, in_edge_fea_len, if_exp=if_exp, normalization=normalization) + self.mp2 = MPLayer(in_atom_fea_len, in_edge_fea_len, in_edge_fea_len, if_exp=if_exp, normalization=normalization) + self.mp3 = MPLayer(in_atom_fea_len, in_edge_fea_len, in_edge_fea_len, if_exp=if_exp, normalization=normalization) + self.mp4 = MPLayer(in_atom_fea_len, in_edge_fea_len, in_edge_fea_len, if_exp=if_exp, normalization=normalization) + self.mp5 = MPLayer(in_atom_fea_len, in_edge_fea_len, in_edge_fea_len - num_l ** 2, if_exp=if_exp, normalization=normalization) + self.lcmp = LCMPLayer(in_atom_fea_len, in_edge_fea_len, num_orbital, num_l, if_exp=if_exp) + self.loss_fn = MaskMSELoss() + + @staticmethod + def _tensor(data, dtype): + if paddle.is_tensor(data): + return paddle.cast(data, dtype) + return paddle.to_tensor(data, dtype=dtype) + + @classmethod + def _prepare_batch(cls, batch): + edge_index = cls._tensor(batch["edge_index"], "int64") + if len(edge_index.shape) == 2 and edge_index.shape[0] != 2: + edge_index = edge_index.transpose([1, 0]) + + prepared = { + "x": cls._tensor(batch["x"], "int64"), + "edge_index": edge_index, + "edge_attr": cls._tensor(batch["edge_attr"], "float32"), + "batch": cls._tensor(batch["batch"], "int64"), + "sub_atom_idx": cls._tensor(batch["sub_atom_idx"], "int64"), + "sub_edge_idx": cls._tensor(batch["sub_edge_idx"], "int64"), + "sub_edge_ang": cls._tensor(batch["sub_edge_ang"], "float32"), + "sub_index": cls._tensor(batch["sub_index"], "int64"), + } + if "label" in batch: + prepared["label"] = cls._tensor(batch["label"], "float32") + if "mask" in batch: + prepared["mask"] = cls._tensor(batch["mask"], "bool") + return prepared + + def forward_backbone(self, atom_attr, edge_idx, edge_attr, batch, sub_atom_idx, sub_edge_idx, sub_edge_ang, sub_index): + atom_attr = paddle.cast(atom_attr, "int64") + edge_idx = paddle.cast(edge_idx, "int64") + batch = paddle.cast(batch, "int64") + distance = edge_attr[:, 0] + + atom_fea0 = self.embed(atom_attr) + edge_fea0 = self.distance_expansion(distance) + + atom_fea, edge_fea = self.mp1(atom_fea0, edge_idx, edge_fea0, batch, distance) + atom_fea, edge_fea = self.mp2(atom_fea, edge_idx, edge_fea, batch, distance) + atom_fea0, edge_fea0 = atom_fea0 + atom_fea, edge_fea0 + edge_fea + atom_fea, edge_fea = self.mp3(atom_fea0, edge_idx, edge_fea0, batch, distance) + atom_fea, edge_fea = self.mp4(atom_fea, edge_idx, edge_fea, batch, distance) + atom_fea0, edge_fea0 = atom_fea0 + atom_fea, edge_fea0 + edge_fea + atom_fea, edge_fea = self.mp5(atom_fea0, edge_idx, edge_fea0, batch, distance) + return self.lcmp(atom_fea, edge_fea, sub_atom_idx, sub_edge_idx, sub_edge_ang, sub_index, distance) + + def forward(self, batch): + batch = self._prepare_batch(batch) + pred = self.forward_backbone( + batch["x"], + batch["edge_index"], + batch["edge_attr"], + batch["batch"], + batch["sub_atom_idx"], + batch["sub_edge_idx"], + batch["sub_edge_ang"], + batch["sub_index"], + ) + label = batch[self.target_name] + mask = batch["mask"] + pred = paddle.reshape(pred, label.shape) + + loss = self.loss_fn(pred, label, mask) + pred_dict = { + self.target_name: pred, + } + loss_dict = { + "loss": loss, + } + return {"loss_dict": loss_dict, "pred_dict": pred_dict} + + def predict(self, batch): + batch = self._prepare_batch(batch) + pred = self.forward_backbone( + batch["x"], + batch["edge_index"], + batch["edge_attr"], + batch["batch"], + batch["sub_atom_idx"], + batch["sub_edge_idx"], + batch["sub_edge_ang"], + batch["sub_index"], + ) + if batch.get("label") is not None: + pred = paddle.reshape(pred, batch["label"].shape) + return { + self.target_name: pred, + } diff --git a/ppmat/models/deeph/deeph_graph_converter.py b/ppmat/models/deeph/deeph_graph_converter.py new file mode 100644 index 00000000..24bd6960 --- /dev/null +++ b/ppmat/models/deeph/deeph_graph_converter.py @@ -0,0 +1,84 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import os + +import numpy as np +from pymatgen.core.structure import Structure +from pymatgen.io.ase import AseAtomsAdaptor + +from ppmat.datasets.deeph_graph import get_graph + + +class DeepHGraphConverter: + """Convert a PaddleMaterials structure graph to a DeepH Hamiltonian graph.""" + + def __init__( + self, + radius, + max_num_nbr, + default_dtype, + interface, + num_l, + create_from_DFT=False, + if_lcmp_graph=True, + separate_onsite=False, + target="hamiltonian", + huge_structure=False, + if_new_sp=False, + ): + self.radius = radius + self.max_num_nbr = max_num_nbr + self.default_dtype = default_dtype + self.interface = interface + self.num_l = num_l + self.create_from_DFT = create_from_DFT + self.if_lcmp_graph = if_lcmp_graph + self.separate_onsite = separate_onsite + self.target = target + self.huge_structure = huge_structure + self.if_new_sp = if_new_sp + + @staticmethod + def _to_structure(structure): + if isinstance(structure, Structure): + return structure + return AseAtomsAdaptor.get_structure(structure) + + def __call__(self, structure, folder, converter_graph=None): + structure = self._to_structure(structure) + default_dtype = self.default_dtype + return get_graph( + np.asarray(structure.cart_coords, dtype=default_dtype), + np.asarray(structure.frac_coords, dtype=default_dtype), + np.asarray(structure.atomic_numbers, dtype=np.int64), + os.path.basename(folder), + r=self.radius, + max_num_nbr=self.max_num_nbr, + numerical_tol=1e-8, + lattice=np.asarray(structure.lattice.matrix, dtype=default_dtype), + default_dtype=default_dtype, + tb_folder=folder, + interface=self.interface, + num_l=self.num_l, + create_from_DFT=self.create_from_DFT, + if_lcmp_graph=self.if_lcmp_graph, + separate_onsite=self.separate_onsite, + target=self.target, + huge_structure=self.huge_structure, + if_new_sp=self.if_new_sp, + converter_graph=converter_graph, + ) diff --git a/ppmat/predictor/__init__.py b/ppmat/predictor/__init__.py index 34a0df0d..1117c3f7 100644 --- a/ppmat/predictor/__init__.py +++ b/ppmat/predictor/__init__.py @@ -13,7 +13,9 @@ # limitations under the License. from ppmat.predictor.base import BasePredictor +from ppmat.predictor.deeph import DeepHPredictor __all__ = [ "BasePredictor", + "DeepHPredictor", ] diff --git a/ppmat/predictor/deeph.py b/ppmat/predictor/deeph.py new file mode 100644 index 00000000..2ff6c4ac --- /dev/null +++ b/ppmat/predictor/deeph.py @@ -0,0 +1,261 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import json +import os +from configparser import ConfigParser +from typing import Dict + +import numpy as np +import paddle +from ase import Atoms +from pymatgen.io.ase import AseAtomsAdaptor + +from ppmat.datasets.deeph_graph import _load_orbital_types +from ppmat.models import build_graph_converter +from ppmat.predictor.base import BasePredictor + + +class DeepHPredictor(BasePredictor): + """Predict DeepH Hamiltonian blocks and export them in DeepH npz format.""" + + @staticmethod + def load_orbital_config(config_path): + config = ConfigParser() + if not config.read(config_path): + raise FileNotFoundError(f"Can not read DeepH config: {config_path}") + return json.loads(config.get("basic", "orbital")) + + @staticmethod + def _load_processed_structure(processed_dir): + lattice = np.loadtxt(os.path.join(processed_dir, "lat.dat")).T + atomic_numbers = ( + np.loadtxt(os.path.join(processed_dir, "element.dat")) + .astype(np.int64) + .reshape(-1) + ) + cart_coords = np.loadtxt(os.path.join(processed_dir, "site_positions.dat")).T + atoms = Atoms( + numbers=atomic_numbers, + positions=cart_coords, + cell=lattice, + pbc=True, + ) + return AseAtomsAdaptor.get_structure(atoms) + + def build_inference_batch( + self, + processed_dir, + species, + num_l, + dtype="float32", + ): + """Build a label-free DeepH batch from overlap-derived local coordinates.""" + structure = self._load_processed_structure(processed_dir) + np_dtype = np.dtype(dtype).type + graph_converter = build_graph_converter( + { + "__class_name__": "DeepHGraphConverter", + "__init_params__": { + "radius": -1.0, + "max_num_nbr": 0, + "default_dtype": np_dtype, + "interface": "npz_rc_only", + "num_l": num_l, + "create_from_DFT": True, + "if_lcmp_graph": True, + "separate_onsite": False, + "target": "hamiltonian", + }, + } + ) + graph = graph_converter(structure, processed_dir) + + species_to_index = { + int(atomic_number): index for index, atomic_number in enumerate(species) + } + atomic_numbers = np.asarray(structure.atomic_numbers, dtype=np.int64) + unknown_species = sorted(set(atomic_numbers) - set(species_to_index)) + if unknown_species: + raise ValueError( + "DeepH inference structure contains species missing from the model: " + f"{unknown_species}" + ) + atom_features = np.asarray( + [species_to_index[int(number)] for number in atomic_numbers], + dtype=np.int64, + ) + subgraph = graph.subgraph_dict + + return { + "x": atom_features, + "edge_index": np.asarray(graph.edge_index, dtype=np.int64), + "edge_attr": np.asarray(graph.edge_attr, dtype=np_dtype), + "batch": np.zeros(atom_features.shape[0], dtype=np.int64), + "pos": np.asarray(structure.cart_coords, dtype=np_dtype), + "sub_atom_idx": np.asarray( + subgraph["subgraph_atom_idx"], dtype=np.int64 + ), + "sub_edge_idx": np.asarray( + subgraph["subgraph_edge_idx"], dtype=np.int64 + ), + "sub_edge_ang": np.asarray( + subgraph["subgraph_edge_ang"], dtype=np_dtype + ), + "sub_index": np.asarray(subgraph["subgraph_index"], dtype=np.int64), + "structure_lattice": np.asarray( + structure.lattice.matrix, dtype=np_dtype + ).reshape(1, 3, 3), + "structure_frac_coords": np.asarray( + structure.frac_coords, dtype=np_dtype + ), + "structure_atomic_numbers": atomic_numbers, + "structure_folder": os.path.abspath(processed_dir), + } + + @staticmethod + def _to_numpy(data): + if paddle.is_tensor(data): + return data.numpy() + return np.asarray(data) + + @staticmethod + def _get_structure_folder(batch) -> str: + folder = batch.get("structure_folder") + if isinstance(folder, (list, tuple)): + if len(folder) != 1: + raise ValueError( + "DeepH Hamiltonian export currently expects batch_size=1." + ) + folder = folder[0] + if not isinstance(folder, str): + raise TypeError("DeepH Hamiltonian export requires `structure_folder`.") + return folder + + def build_hamiltonian_npz( + self, + pred, + batch, + orbital, + spinful: bool = False, + ) -> Dict[str, np.ndarray]: + """Convert DeepH edge predictions to rh.npz-compatible blocks.""" + if spinful: + raise NotImplementedError( + "Spinful DeepH Hamiltonian export is not enabled." + ) + + folder = self._get_structure_folder(batch) + atom_num_orbital = _load_orbital_types( + os.path.join(folder, "orbital_types.dat") + ) + + pred = self._to_numpy(pred) + edge_index = self._to_numpy(batch["edge_index"]).astype(np.int64) + if edge_index.ndim == 2 and edge_index.shape[0] != 2: + edge_index = edge_index.T + edge_attr = self._to_numpy(batch["edge_attr"]) + atomic_numbers = self._to_numpy(batch["structure_atomic_numbers"]).astype( + np.int64 + ) + lattice = self._to_numpy(batch["structure_lattice"]).reshape(-1, 3, 3)[0] + inv_lattice = np.linalg.inv(lattice) + + hamiltonian = {} + lattice_shifts = np.rint( + edge_attr[:, 4:7] @ inv_lattice - edge_attr[:, 7:10] @ inv_lattice + ).astype(int) + + for edge_id in range(edge_attr.shape[0]): + atom_i, atom_j = edge_index[:, edge_id] + block = np.zeros( + (atom_num_orbital[atom_i], atom_num_orbital[atom_j]), + dtype=pred.dtype, + ) + atomic_number_i = int(atomic_numbers[atom_i]) + atomic_number_j = int(atomic_numbers[atom_j]) + + for out_idx, orbital_dict in enumerate(orbital): + for n_m_str, orbital_pair in orbital_dict.items(): + condition_i, condition_j = map(int, n_m_str.split()) + if ( + atomic_number_i == condition_i + and atomic_number_j == condition_j + ): + orbital_i, orbital_j = orbital_pair + block[orbital_i, orbital_j] = pred[edge_id, out_idx] + + key = json.dumps( + [ + int(lattice_shifts[edge_id, 0]), + int(lattice_shifts[edge_id, 1]), + int(lattice_shifts[edge_id, 2]), + int(atom_i) + 1, + int(atom_j) + 1, + ] + ) + hamiltonian[key] = block + return hamiltonian + + def save_hamiltonian_npz( + self, + path, + pred, + batch, + orbital, + spinful: bool = False, + ) -> None: + """Save DeepH predictions as an rh.npz-compatible Hamiltonian file.""" + hamiltonian = self.build_hamiltonian_npz( + pred=pred, + batch=batch, + orbital=orbital, + spinful=spinful, + ) + os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) + np.savez(path, **hamiltonian) + + def predict_hamiltonian( + self, + processed_dir, + output_path, + orbital, + species, + num_l, + dtype="float32", + ): + """Run Paddle DeepH inference and save predicted Hamiltonian blocks.""" + batch = self.build_inference_batch( + processed_dir=processed_dir, + species=species, + num_l=num_l, + dtype=dtype, + ) + if self.eval_with_no_grad: + with paddle.no_grad(): + prediction = self.model.predict(batch) + else: + prediction = self.model.predict(batch) + prediction = self.post_process(prediction) + pred = prediction[self.model.target_name] + self.save_hamiltonian_npz( + output_path, + pred=pred, + batch=batch, + orbital=orbital, + spinful=False, + ) + return output_path diff --git a/ppmatSim/configs/Calculator/deeph_openmx.yaml b/ppmatSim/configs/Calculator/deeph_openmx.yaml new file mode 100644 index 00000000..a1150819 --- /dev/null +++ b/ppmatSim/configs/Calculator/deeph_openmx.yaml @@ -0,0 +1,10 @@ +_target_: ppmat.calculator.deeph.DeepHDFTCalculator +type: deeph_dft +calculator_cls: ase.calculators.openmx.OpenMX +calculator_kwargs: + xc: PBE + scf_kgrid: [1, 1, 1] + scf_energycutoff: 200 +command: null +output_filename: openmx.log +overlap_pattern: output/overlaps_*.h5 diff --git a/ppmatSim/configs/Task/deeph_dft_inference.yaml b/ppmatSim/configs/Task/deeph_dft_inference.yaml new file mode 100644 index 00000000..529e7b75 --- /dev/null +++ b/ppmatSim/configs/Task/deeph_dft_inference.yaml @@ -0,0 +1,9 @@ +_target_: ppmat.calculator.deeph.DeepHDFTInferenceTask +orbital: null +orbital_config_path: null +species: [6] +num_l: 5 +output_dir: deeph_dft_results +dtype: float32 +run_dft: true +processed_dirs: null diff --git a/ppmatSim/configs/deeph_openmx.yaml b/ppmatSim/configs/deeph_openmx.yaml new file mode 100644 index 00000000..4b679adf --- /dev/null +++ b/ppmatSim/configs/deeph_openmx.yaml @@ -0,0 +1,22 @@ +# @package _global_ +device: gpu + +defaults: + - /Run: default + - /Logger: default + - /Model: load_model + - /System: load_system + - /Calculator: deeph_openmx + - /Task: deeph_dft_inference + - override hydra/job_logging: disabled + - _self_ + +hydra: + job: + name: deeph_openmx + chdir: True + run: + dir: ./out_${hydra:job.name}/${now:%Y%m%d_%H%M%S} + +Predictor: + _target_: ppmat.predictor.DeepHPredictor diff --git a/ppmatSim/main.py b/ppmatSim/main.py index dac5abfb..8f230228 100644 --- a/ppmatSim/main.py +++ b/ppmatSim/main.py @@ -21,7 +21,6 @@ from omegaconf import DictConfig from omegaconf import OmegaConf -from ppmat.predictor import BasePredictor from ppmat.predictor.structures import build_init_structures from ppmat.utils import logger @@ -40,9 +39,19 @@ def main(cfg: DictConfig): # Initialize the model load_model = instantiate(cfg.Model) - predictor = BasePredictor( - work_dir=cfg.Run.work_dir, device=cfg.device, **load_model - ) + if cfg.get("Predictor") is None: + from ppmat.predictor import BasePredictor + + predictor = BasePredictor( + work_dir=cfg.Run.work_dir, device=cfg.device, **load_model + ) + else: + predictor = instantiate( + cfg.Predictor, + work_dir=cfg.Run.work_dir, + device=cfg.device, + **load_model, + ) # Detect interface type and interface object if cfg.get("Calculator") is not None: diff --git a/requirements.txt b/requirements.txt index 0c8da8a9..324a0ead 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,6 +2,7 @@ ase==3.23.0 colorlog==6.8.2 Cython==3.0.12 hydra-core==1.3.2 +h5py==3.12.1 einops==0.8.0 importlib_metadata==8.6.1 jarvis_tools==2025.5.30 diff --git a/test/test_ase_dft.py b/test/test_ase_dft.py new file mode 100644 index 00000000..0b5f9a16 --- /dev/null +++ b/test/test_ase_dft.py @@ -0,0 +1,133 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +import tempfile +import unittest +from types import SimpleNamespace + +import h5py +import numpy as np +from ase import Atoms + +from ppmat.calculator.deeph import DeepHDFTCalculator +from ppmat.calculator.deeph import preprocess_openmx_overlap +from ppmat.models.deeph import DeepHHamiltonian +from ppmat.predictor.deeph import DeepHPredictor + + +class TestASEDFT(unittest.TestCase): + @staticmethod + def _write_openmx_overlap_fixture(raw_dir): + os.makedirs(os.path.join(raw_dir, "output"), exist_ok=True) + with open(os.path.join(raw_dir, "openmx.out"), "w") as file: + file.write( + """ +Atoms.UnitVectors.Unit Ang + +Fractional coordinates of the final structure +index element x y z +--- +--- +1 C 0.0 0.0 0.0 +2 C 0.25 0.0 0.0 +3 C 0.0 0.25 0.0 + +""" + ) + overlap_path = os.path.join(raw_dir, "output", "overlaps_0.h5") + with h5py.File(overlap_path, "w") as overlap_file: + for atom_i in range(1, 4): + for atom_j in range(1, 4): + key = json.dumps([0, 0, 0, atom_i, atom_j]) + overlap_file[key] = np.eye(4, dtype=np.float64) + + def test_ase_calculator_schedules_backend(self): + atoms = Atoms( + symbols=["Cu"], + positions=[[0.0, 0.0, 0.0]], + cell=np.eye(3) * 4.0, + pbc=True, + ) + with tempfile.TemporaryDirectory() as tmpdir: + calculator = DeepHDFTCalculator( + predictor=SimpleNamespace(), + calculator_cls="ase.calculators.emt.EMT", + ) + calculator._run_dft(atoms, tmpdir) + self.assertTrue(os.path.exists(os.path.join(tmpdir, "structure.xyz"))) + + def test_openmx_overlap_to_deeph_inference(self): + with tempfile.TemporaryDirectory() as tmpdir: + raw_dir = os.path.join(tmpdir, "raw") + processed_dir = os.path.join(tmpdir, "processed") + self._write_openmx_overlap_fixture(raw_dir) + + preprocess_openmx_overlap(raw_dir, processed_dir) + predictor = DeepHPredictor() + batch = predictor.build_inference_batch( + processed_dir, + species=[6], + num_l=2, + ) + self.assertEqual(batch["edge_index"].shape, (2, 9)) + self.assertEqual(batch["sub_edge_ang"].shape[-1], 4) + + orbital = [ + {"6 6": [orbital_i, orbital_j]} + for orbital_i in range(4) + for orbital_j in range(4) + ] + + model = DeepHHamiltonian( + num_species=1, + in_atom_fea_len=16, + in_edge_fea_len=32, + num_orbital=16, + num_l=2, + gauss_stop=6.0, + if_exp=True, + normalization="LayerNorm", + target_name="label", + ) + predictor.model = model + predictor.eval_with_no_grad = True + predictor.post_process = lambda data: data + calculator = DeepHDFTCalculator(predictor=predictor) + calculator.run_deeph_inference( + structures=[Atoms("C3")], + orbital=orbital, + species=[6], + num_l=2, + output_dir=os.path.join(tmpdir, "result"), + dtype="float32", + run_dft=False, + processed_dirs=[processed_dir], + ) + + output_path = os.path.join(tmpdir, "result", "000000", "rh_pred.npz") + with np.load(output_path) as prediction: + self.assertEqual(len(prediction.files), 9) + self.assertEqual(prediction[prediction.files[0]].shape, (4, 4)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_deeph.py b/test/test_deeph.py new file mode 100644 index 00000000..7522e2b1 --- /dev/null +++ b/test/test_deeph.py @@ -0,0 +1,176 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import tempfile +import unittest + +import numpy as np +import paddle +from ase import Atoms +from omegaconf import OmegaConf +from pymatgen.core.structure import Structure + +from ppmat.datasets.collate_fn import DefaultCollator +from ppmat.datasets.custom_data_type import ConcatData +from ppmat.models.deeph import DeepHGraphConverter +from ppmat.models.deeph import DeepHHamiltonian +from ppmat.predictor.deeph import DeepHPredictor + + +class TestDeepH(unittest.TestCase): + def test_deeph_config_can_be_loaded(self): + config_path = os.path.join( + "electronic_structure", + "configs", + "deeph", + "deeph_graphene.yaml", + ) + config = OmegaConf.to_container(OmegaConf.load(config_path), resolve=True) + + self.assertEqual(config["Model"]["__class_name__"], "DeepHHamiltonian") + self.assertEqual( + config["Dataset"]["train"]["dataset"]["__class_name__"], + "DeepHDataset", + ) + init_params = config["Dataset"]["train"]["dataset"]["__init_params__"] + self.assertIn("config", init_params) + self.assertNotIn("config_files", init_params) + self.assertEqual(len(init_params["config"]["basic"]["orbital"]), 169) + + def test_deeph_forward(self): + paddle.seed(42) + model = DeepHHamiltonian( + num_species=1, + in_atom_fea_len=16, + in_edge_fea_len=32, + num_orbital=9, + num_l=3, + gauss_stop=6.0, + if_exp=True, + normalization="LayerNorm", + target_name="label", + ) + + num_nodes = 3 + num_edges = 4 + num_sub_edges = 8 + batch = { + "x": paddle.zeros([num_nodes], dtype="int64"), + "edge_index": paddle.to_tensor([[0, 1, 1, 2], [1, 0, 2, 1]], dtype="int64"), + "edge_attr": paddle.ones([num_edges, 1], dtype="float32"), + "batch": paddle.zeros([num_nodes], dtype="int64"), + "sub_atom_idx": paddle.to_tensor( + [ + [0, 1], + [1, 0], + [1, 2], + [2, 1], + [0, 1], + [1, 0], + [1, 2], + [2, 1], + ], + dtype="int64", + ), + "sub_edge_idx": paddle.to_tensor([0, 1, 2, 3, 0, 1, 2, 3], dtype="int64"), + "sub_edge_ang": paddle.zeros([num_sub_edges, 9], dtype="float32"), + "sub_index": paddle.to_tensor([0, 1, 2, 3, 4, 5, 6, 7], dtype="int64"), + "label": paddle.zeros([num_edges, 9], dtype="float32"), + "mask": paddle.ones([num_edges, 9], dtype="bool"), + } + + output = model(batch) + + self.assertIn("loss", output["loss_dict"]) + self.assertIn("label", output["pred_dict"]) + self.assertEqual(list(output["pred_dict"]["label"].shape), [num_edges, 9]) + + def test_deeph_data_uses_default_collator(self): + sample = { + "x": ConcatData(np.zeros([2], dtype=np.int64)), + "edge_index": ConcatData(np.asarray([[0, 1], [1, 0]], dtype=np.int64)), + "edge_attr": ConcatData(np.ones([2, 1], dtype=np.float32)), + "batch": ConcatData(np.zeros([2], dtype=np.int64)), + "label": ConcatData(np.zeros([2, 9], dtype=np.float32)), + "mask": ConcatData(np.ones([2, 9], dtype=bool)), + "sub_atom_idx": ConcatData(np.asarray([[0, 1], [1, 0]], dtype=np.int64)), + "sub_edge_idx": ConcatData(np.asarray([0, 1], dtype=np.int64)), + "sub_edge_ang": ConcatData(np.zeros([2, 9], dtype=np.float32)), + "sub_index": ConcatData(np.asarray([0, 1], dtype=np.int64)), + } + + collated = DefaultCollator()([sample]) + + self.assertIsInstance(collated, dict) + self.assertEqual(collated["x"].tolist(), [0, 0]) + self.assertEqual(collated["sub_atom_idx"].tolist(), [[0, 1], [1, 0]]) + self.assertEqual(collated["sub_edge_idx"].tolist(), [0, 1]) + self.assertEqual(collated["sub_index"].tolist(), [0, 1]) + + def test_deeph_prediction_can_export_hamiltonian_npz(self): + with tempfile.TemporaryDirectory() as tmpdir: + with open(os.path.join(tmpdir, "orbital_types.dat"), "w") as f: + f.write("0 1\n") + f.write("0 1\n") + + batch = { + "edge_index": np.asarray([[0], [1]], dtype=np.int64), + "edge_attr": np.asarray( + [[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0]], + dtype=np.float32, + ), + "structure_lattice": np.eye(3, dtype=np.float32).reshape(1, 3, 3), + "structure_atomic_numbers": np.asarray([6, 6], dtype=np.int64), + "structure_folder": tmpdir, + } + pred = np.asarray([[0.25]], dtype=np.float32) + orbital = [{"6 6": [0, 1]}] + predictor = DeepHPredictor() + + hamiltonian = predictor.build_hamiltonian_npz(pred, batch, orbital) + + self.assertEqual(list(hamiltonian.keys()), ["[0, 0, 0, 1, 2]"]) + self.assertEqual(hamiltonian["[0, 0, 0, 1, 2]"].shape, (4, 4)) + self.assertAlmostEqual( + float(hamiltonian["[0, 0, 0, 1, 2]"][0, 1]), + 0.25, + places=6, + ) + + output_path = os.path.join(tmpdir, "rh_pred.npz") + predictor.save_hamiltonian_npz(output_path, pred, batch, orbital) + loaded = np.load(output_path) + self.assertIn("[0, 0, 0, 1, 2]", loaded.files) + self.assertAlmostEqual( + float(loaded["[0, 0, 0, 1, 2]"][0, 1]), + 0.25, + places=6, + ) + + def test_deeph_graph_converter_accepts_ase_atoms(self): + atoms = Atoms( + symbols=["C", "C"], + positions=[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + cell=np.eye(3) * 4.0, + pbc=True, + ) + structure = DeepHGraphConverter._to_structure(atoms) + + self.assertIsInstance(structure, Structure) + self.assertEqual(len(structure), 2) + + +if __name__ == "__main__": + unittest.main()