diff --git a/Dockerfile b/Dockerfile index 7d0d7c46..46a39c35 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,57 +1,193 @@ -# syntax=docker/dockerfile:1 -FROM python:3.10-bookworm - -SHELL ["/bin/bash", "-c"] -WORKDIR /src/temp - -#Install Featomic -RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs > installRust.sh \ - && chmod 700 installRust.sh \ - && ./installRust.sh -y \ - && source $HOME/.cargo/env \ - && pip install git+https://github.com/metatensor/featomic.git - -#Install OpenMP -RUN wget https://download.open-mpi.org/release/open-mpi/v5.0/openmpi-5.0.2.tar.bz2 \ - && tar -xf openmpi-5.0.2.tar.bz2 \ - && cd openmpi-5.0.2 \ - && ./configure --prefix=/usr/local/openmpi5/lib \ - && make all install \ - && cd .. \ - && rm -R openmpi-5.0.2 - -ENV PATH=/usr/local/openmpi5/lib/bin:$PATH -ENV LD_LIBRARY_PATH=/usr/local/openmpi5/lib -ENV HDF5_DIR=/usr/local/hdf5 - -#Install HDF5 -RUN wget https://hdf-wordpress-1.s3.amazonaws.com/wp-content/uploads/manual/HDF5/HDF5_1_14_3/src/hdf5-1.14.3.tar.gz \ - && tar -xf hdf5-1.14.3.tar.gz \ - && cd hdf5-1.14.3 \ - && HDF5_MPI="ON" CC=mpicc ./configure --enable-shared --enable-parallel --prefix=/usr/local/hdf5 \ - && HDF5_DIR=/usr/local/hdf5 make \ - && HDF5_DIR=/usr/local/hdf5 make install \ - && cd .. \ - && rm -R hdf5-1.14.3 - -#Install h5py -RUN HDF5_DIR=/usr/local/hdf5 HDF5_MPI="ON" CC=mpicc pip install --no-cache-dir --no-binary=h5py h5py - -RUN apt-get update && apt-get install -y \ - ninja-build \ - gfortran \ - && rm -rf /var/lib/apt/lists/* - -RUN pip install meson cython \ - && pip install --prefer-binary pyscf - -#Install SALTED -COPY . /src/temp/SALTED-master -RUN cd /src/temp/SALTED-master \ - && make \ - && pip install . - -RUN rm -R /src/temp +# syntax=docker/dockerfile:1.7 + +ARG PYTHON_IMAGE=python:3.10-slim-bookworm +ARG OPENMPI_VERSION=4.1.8 +ARG HDF5_VERSION=1.14.3 + + +# ----------------------------------------------------------------------------- +# Build MPI and parallel HDF5 +# ----------------------------------------------------------------------------- +FROM ${PYTHON_IMAGE} AS native-builder + +ARG OPENMPI_VERSION +ARG HDF5_VERSION + +ENV DEBIAN_FRONTEND=noninteractive + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \ + --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \ + apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + curl \ + gfortran \ + libevent-dev \ + libhwloc-dev \ + libmunge-dev \ + libpmi2-0-dev \ + zlib1g-dev + +WORKDIR /tmp/build + +# Open MPI +# +# Use the PMIx version bundled with Open MPI instead of independently +# combining Open MPI 4.1 with PMIx 6. +RUN curl -fsSL \ + "https://download.open-mpi.org/release/open-mpi/v4.1/openmpi-${OPENMPI_VERSION}.tar.gz" \ + | tar -xz \ + && cd "openmpi-${OPENMPI_VERSION}" \ + && ./configure \ + --prefix=/opt/mpi \ + --enable-shared \ + --disable-static \ + --disable-debug \ + --enable-builtin-atomics \ + --disable-mpi-fortran \ + --disable-oshmem \ + --with-slurm \ + --with-pmix=internal \ + --with-hwloc=/usr \ + --with-libevent=/usr \ + --with-zlib=/usr \ + --without-psm \ + --without-psm2 \ + && make -j"$(nproc)" \ + && make install-strip \ + && cd /tmp/build \ + && rm -rf "openmpi-${OPENMPI_VERSION}" + +ENV PATH=/opt/mpi/bin:${PATH} +ENV LD_LIBRARY_PATH=/opt/mpi/lib + +# Parallel HDF5 +RUN curl -fsSL \ + "https://hdf-wordpress-1.s3.amazonaws.com/wp-content/uploads/manual/HDF5/HDF5_1_14_3/src/hdf5-${HDF5_VERSION}.tar.gz" \ + | tar -xz \ + && cd "hdf5-${HDF5_VERSION}" \ + && CC=/opt/mpi/bin/mpicc ./configure \ + --prefix=/opt/hdf5 \ + --enable-shared \ + --disable-static \ + --enable-parallel \ + && make -j"$(nproc)" \ + && make install-strip \ + && cd /tmp/build \ + && rm -rf "hdf5-${HDF5_VERSION}" + +# Create a runtime-only copy without headers, pkg-config files, static +# archives, documentation, or HDF5 developer tools. +RUN mkdir -p /opt/runtime \ + && cp -a /opt/mpi /opt/runtime/mpi \ + && cp -a /opt/hdf5 /opt/runtime/hdf5 \ + && rm -rf \ + /opt/runtime/mpi/include \ + /opt/runtime/mpi/share/man \ + /opt/runtime/mpi/share/doc \ + /opt/runtime/mpi/lib/pkgconfig \ + /opt/runtime/hdf5/include \ + /opt/runtime/hdf5/share \ + /opt/runtime/hdf5/bin \ + /opt/runtime/hdf5/lib/pkgconfig \ + && find /opt/runtime -type f \ + \( -name '*.a' -o -name '*.la' \) \ + -delete + + +# ----------------------------------------------------------------------------- +# Build the Python environment +# ----------------------------------------------------------------------------- +FROM native-builder AS python-builder + +ENV HDF5_DIR=/opt/hdf5 +ENV PATH=/opt/venv/bin:/opt/mpi/bin:${PATH} +ENV LD_LIBRARY_PATH=/opt/mpi/lib:/opt/hdf5/lib +ENV PIP_DISABLE_PIP_VERSION_CHECK=1 + +RUN python -m venv /opt/venv + +RUN --mount=type=cache,target=/root/.cache/pip \ + python -m pip install --upgrade \ + pip \ + setuptools \ + wheel \ + build \ + && python -m pip install \ + featomic \ + numpy \ + cython \ + pkgconfig \ + && MPICC=/opt/mpi/bin/mpicc \ + python -m pip install \ + --no-binary=mpi4py \ + mpi4py \ + && HDF5_DIR=/opt/hdf5 \ + HDF5_MPI=ON \ + CC=/opt/mpi/bin/mpicc \ + python -m pip install \ + --no-build-isolation \ + --no-binary=h5py \ + h5py \ + && python -m pip install \ + meson \ + packaging \ + numba \ + ase \ + scipy \ + pyyaml \ + sympy \ + && python -m pip install \ + --prefer-binary \ + pyscf + +# Copy SALTED last so that source-code changes do not invalidate the expensive +# MPI, HDF5, and Python dependency layers. +WORKDIR /src/SALTED +COPY . . + +RUN --mount=type=cache,target=/root/.cache/pip \ + python -m pip install . \ + && python -m pip uninstall -y \ + build \ + cython \ + meson \ + pkgconfig \ + wheel \ + && find /opt/venv -type d -name '__pycache__' \ + -prune -exec rm -rf '{}' + + + +# ----------------------------------------------------------------------------- +# Minimal runtime +# ----------------------------------------------------------------------------- +FROM ${PYTHON_IMAGE} AS runtime + +ENV DEBIAN_FRONTEND=noninteractive + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \ + --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \ + apt-get update && apt-get install -y --no-install-recommends \ + ca-certificates \ + libevent-2.1-7 \ + libevent-pthreads-2.1-7 \ + libgfortran5 \ + libhwloc15 \ + libmunge2 \ + libpmi2-0 \ + openssh-client \ + zlib1g + +COPY --from=native-builder /opt/runtime/mpi /opt/mpi +COPY --from=native-builder /opt/runtime/hdf5 /opt/hdf5 +COPY --from=python-builder /opt/venv /opt/venv + +ENV PATH=/opt/venv/bin:/opt/mpi/bin:${PATH} +ENV LD_LIBRARY_PATH=/opt/mpi/lib:/opt/hdf5/lib +ENV HDF5_DIR=/opt/hdf5 +ENV PYTHONDONTWRITEBYTECODE=1 +ENV PYTHONUNBUFFERED=1 WORKDIR /work -ENTRYPOINT ["/bin/bash"] + +CMD ["/bin/bash"] diff --git a/docs/binary_models.md b/docs/binary_models.md new file mode 100644 index 00000000..e826ebc1 --- /dev/null +++ b/docs/binary_models.md @@ -0,0 +1,107 @@ +# Binary Salted Models +By running `salted_pack.py`, a single binary file is created that consolidates and serializes multiple data sources (NumPy arrays, HDF5 datasets, and model parameters) into one portable format. +This simplifies model sharing and enables easier deployment for prediction tasks. + +## Usage +### Creating binary models +Only a single command is required to transform a newly trained and verified model into the binary `.salted` format. +The command must be executed from the main project directory: + + salted_pack.py + +This will generate a `.salted` file. + +### Deploying binary models +Predictions using a binary model can be performed with: + + predict_from_model.py + +**Input arguments** +| Argument | Meaning | +| -------- | --------| +| model | Path to the `.salted` model file. | +| xyz | Path to the `.xyz` file containing the structure for which the prediction is performed. | +| -o | Output directory for the predicted coefficients (optional, default: `predictions/`) | + +## Implementation details +### Basic design +All integers are little-endian signed 32-bit unless noted; all floating arrays are little-endian float64. All 5-char names are ASCII padded with NUL to 5 bytes. +General Arrays + +In the following every array is encoded using the same scheme: + + def encode_array(NDIMS:int, DIMS:list[int], data): + file.write(NDIMS) + for dim_size in DIMS: + file.write(dim_size) + file.write(data) + +This results in a datastructure as follows: + + NDIMS (int32) + DIMS (int32 * NDIMS) + DATA (TYPE OF ARRAY) + +The Type of the given array is encoded using the numbers from 0 - 5: + + int32=0 + int64=1 + float32=2 + float64=3 + str=4 + bool=5 + +### Sections in the file +Now the different sections of the file are explained in more detail: +#### Container header + + MAGIC (5 bytes): b"SALTD" + VERSION (int32) + N_BLOCKS (int32) + TOC: N_BLOCKS entries of: + BLOCK_NAME (5 bytes, NUL-padded) + BLOCK_OFFSET (int32): file offset (from start) where the block payload begins + +#### AVERG, WIG, FPS, WEIGH + + TYPE (int32) (datatype of the following block) + NFILES (int32) (number of arrays in the specific key) + FOR EACH FILE + encode_array(NDIMS, DIMS, data) + +#### FEATS, PROJE + + TYPE (int32): float64 + NKEYS (int32) — top-level HDF5 group count (sorted) + For each top-level key: + KEY5 (5 bytes, NUL-padded) + NSUB (int32) — number of datasets under this key (sorted) + For each sub-key dataset: + encode_array(NDIMS, DIMS, data) + +#### CONFG + + For each entry in inputs (fixed order in code): + KEY5 (5 bytes) — e.g., b"averg", b"ncut\0", … + TYPE (int32) + VALUE encoded by VAL_TYPE: + bool: int32 (0 or 1) + int32: int32 + float64: float64 + str: SLEN (int32, byte length), then SLEN bytes UTF-8 + +#### BASIS (Only if pyscf is installed) + + TYPE (int32): float64 (tag for numeric arrays in this block) + NELEM (int32): number of elements included + For each element: + ELEM_ID (int32): PySCF element index + Four arrays follow, each preceded by a shape header: + contractions_per_shell (int32[]) + encode_array(NDIMS, DIMS, data) + angular_momenta_per_shell (int32[]) + encode_array(NDIMS, DIMS, data) + exponents_per_shell (float64[]) + encode_array(NDIMS, DIMS, data) + coeffs_per_shell (float64[]) + encode_array(NDIMS, DIMS, data) diff --git a/docs/docker.md b/docs/docker.md index b290eb71..28b9a6a6 100644 --- a/docs/docker.md +++ b/docs/docker.md @@ -21,6 +21,7 @@ docker build -f Dockerfile -t salted . ``` Alternatively, to build using Podman, run: + ```bash podman build --format docker -f Dockerfile -t salted . ``` @@ -28,38 +29,40 @@ podman build --format docker -f Dockerfile -t salted . This will produce a local container image named `salted`. ## Building for HPC Cluster (Apptainer) + For use on clusters managed by Slurm, the image must be converted into an Apptainer (`.sif`) format. Ensure **Apptainer** is available on all nodes. ### Workflow 1. Build the container image (using Docker or Podman): - ```bash - podman build --format docker -f Dockerfile -t salted . - ``` -2. Create a container instance from the image: ```bash - podman create --name=salted --hostname=salted salted:latest - ``` -3. Export the container filesystem: -```bash - mkdir salted - docker export salted | tar -C salted -xf - - ``` -4. Generate a runtime configuration: _(If runc is unavailable, crun can be used instead.)_ + podman build --format docker -t salted:latest . +``` + +2. Create a tarball of the container image: ```bash - cd salted - runc spec --rootless - cd .. - ``` -5. Build the Apptainer image: + podman save -o salted.tar salted:latest +``` + +3. Build the Apptainer image from the tarball[^1]: ```bash - apptainer build salted.sif salted - ``` + apptainer build salted.sif docker-archive://$(pwd)/salted.tar +``` +[^1]: It is very important to use the `docker-archive` URI scheme to ensure proper handling of the image format. Do not simply pass the tarball path directly to `apptainer build`, as this will lead to errors. ### Running the Container + To execute a command within the Apptainer container: + ```bash apptainer exec salted.sif [COMMAND] - ``` +``` + +To run a parallel job usig Slurm, use `srun` with the Apptainer image: + +```bash + srun --ntasks=4 --mpi=pmi2 apptainer exec salted.sif [COMMAND] +``` + This setup allows seamless integration of SALTED across local development environments and HPC systems, maintaining consistent dependencies and runtime behavior. diff --git a/mkdocs.yaml b/mkdocs.yaml index 44f2a9dd..72f27b3d 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -11,8 +11,10 @@ site_url: https://github.com/andreagrisafi/SALTED nav: - Home: index.md - Installation: installation.md + - Docker: docker.md - Theory: theory.md - Workflow: workflow.md + - Binary Deployment: binary_models.md - Input: input.md - Tutorial: - Part 1 - Dataset: tutorial/dataset.md diff --git a/salted/minimize_loss.py b/salted/minimize_loss.py index 9b6cfa4f..6e4aa6e6 100644 --- a/salted/minimize_loss.py +++ b/salted/minimize_loss.py @@ -137,7 +137,7 @@ def build(): trainrange = trainrangetot[:ntraintot] ntrain = int(len(trainrange)) - def loss_func(weights, ovlp_list, psi_list): + def loss_func(weights, ovlp_list, psi_list, coef_list): """Given the weight-vector of the RKHS, compute the gradient of the electron-density loss function.""" # global totsize @@ -152,13 +152,7 @@ def loss_func(weights, ovlp_list, psi_list): # loop over training structures for iconf in range(ntrain): - ref_coefs = np.load( - osp.join( - saltedpath, - "coefficients", - f"coefficients_conf{trainrange[iconf]}.npy", - ) - ) + ref_coefs = coef_list[iconf] if average: Av_coeffs = np.zeros(ref_coefs.shape[0]) @@ -227,7 +221,7 @@ def loss_func(weights, ovlp_list, psi_list): return loss - def grad_func(weights, ovlp_list, psi_list): + def grad_func(weights, ovlp_list, psi_list, coef_list): """ Given the weight-vector of the RKHS, compute the gradient of the electron-density loss function. """ @@ -243,10 +237,7 @@ def grad_func(weights, ovlp_list, psi_list): # loop over training structures for iconf in range(ntrain): - # load reference QM data - ref_coefs = np.load(osp.join( - saltedpath, "coefficients", f"coefficients_conf{trainrange[iconf]}.npy" - )) + ref_coefs = coef_list[iconf] if average: Av_coeffs = np.zeros(ref_coefs.shape[0]) @@ -363,6 +354,7 @@ def curv_func(cg_dire, ovlp_list, psi_list): print("loading matrices...") ovlp_list = [] psi_list = [] + coef_list = [] for iconf in trainrange: ovlp_list.append( np.load(osp.join(saltedpath, "overlaps", f"overlap_conf{iconf}.npy")) @@ -372,6 +364,9 @@ def curv_func(cg_dire, ovlp_list, psi_list): psi_list.append(sparse.load_npz(osp.join( saltedpath, fdir, f"M{Menv}_zeta{zeta}", f"psi-nm_conf{iconf}.npz" ))) + coef_list.append(np.load(osp.join( + saltedpath, "coefficients", f"coefficients_conf{iconf}.npy" + ))) elif saltedtype=="density-response": for icart in ["x","y","z"]: psi_list.append(sparse.load_npz(osp.join( @@ -418,7 +413,7 @@ def curv_func(cg_dire, ovlp_list, psi_list): r = np.load(rpath) s = np.multiply(P, r) delnew = np.dot(r, s) - loss = loss_func(w, ovlp_list, psi_list) + loss = loss_func(w, ovlp_list, psi_list, coef_list) else: # Print a warning and revert to the else behavior print( @@ -427,8 +422,8 @@ def curv_func(cg_dire, ovlp_list, psi_list): if init: w = np.ones(totsize) * 1e-04 - loss = loss_func(w, ovlp_list, psi_list) - r = -grad_func(w, ovlp_list, psi_list) + loss = loss_func(w, ovlp_list, psi_list, coef_list) + r = -grad_func(w, ovlp_list, psi_list, coef_list) d = np.multiply(P, r) delnew = np.dot(r, d) @@ -470,11 +465,11 @@ def curv_func(cg_dire, ovlp_list, psi_list): ) if (i+1)%50==0: loss_old = loss.copy() - loss = loss_func(w, ovlp_list, psi_list) + loss = loss_func(w, ovlp_list, psi_list, coef_list) if loss>loss_old: if rank == 0: print(f"WARNING: loss function increased, search direction reset as the steepest descent.") - r = -grad_func(w, ovlp_list, psi_list) + r = -grad_func(w, ovlp_list, psi_list, coef_list) if rank == 0: print(f"step {i+1}, gradient norm: {np.linalg.norm(r):.3e}, loss: {loss:.3e}", flush=True) if np.linalg.norm(r) < gradtol: diff --git a/salted/pack_model.py b/salted/pack_model.py new file mode 100644 index 00000000..d914c130 --- /dev/null +++ b/salted/pack_model.py @@ -0,0 +1,430 @@ +import h5py +import glob, os +import numpy as np +import argparse + +import struct + +def u32(x): return struct.pack(' 5: + raise ValueError(f"Key longer than 5 bytes: {s!r}") + f.write(s + b'\0' * (5 - len(s))) + +def write_data_head(SALTED_file, data): + SALTED_file.write(i32(int(data.ndim))) + for dim in data.shape: + SALTED_file.write(i32(int(dim))) + + +def write_chunk_location(SALTED_file, chunk_name, location): + global OFFSET_TABLE_OF_CONTENTS + end_of_block = SALTED_file.tell() + SALTED_file.seek(OFFSET_TABLE_OF_CONTENTS) + write_key5(SALTED_file, chunk_name.encode("utf-8")) + SALTED_file.write(i32(int(location))) + SALTED_file.seek(end_of_block) + OFFSET_TABLE_OF_CONTENTS += 5+4 + + + +#Format: +#TYPE_OF_DATA (4 bytes, int32) +#nfiles (4 bytes, int32) +#FOR EACH FILE: + #element (5 bytes, str) + #ndims (4 bytes, int32) + #dims (4*ndims bytes, int32) + #data (ndims*8 bytes, float64) +def pack_averages(SALTED_file, path, debug: bool = False): + if debug: print("Writing Averages") + begin_of_block = SALTED_file.tell() + SALTED_file.write(i32(int(types_dict["float64"]))) + files = glob.glob(os.path.join(path,"coefficients","averages",'av*.npy')) + SALTED_file.write(i32(int(len(files)))) + for file in files: + if debug: print(file) + element = os.path.basename(file).split('.')[0].split("_")[-1] + write_key5(SALTED_file, element.encode("utf-8")) + data = np.load(file).astype(np.float64) + write_data_head(SALTED_file, data) + SALTED_file.write(np.asarray(data,dtype=' 0), + (b"ncut\0", inp.descriptor.sparsify.ncut), #int32 + (b"nang1", inp.descriptor.rep1.nang), + (b"nang2", inp.descriptor.rep2.nang), + (b"nrad1", inp.descriptor.rep1.nrad), + (b"nrad2", inp.descriptor.rep2.nrad), + (b"Menv\0", inp.gpr.Menv), + (b"Ntran", inp.gpr.Ntrain), + (b"rcut1", inp.descriptor.rep1.rcut), #float64 + (b"rcut2", inp.descriptor.rep2.rcut), + (b"sig1\0", inp.descriptor.rep1.sig), + (b"sig2\0", inp.descriptor.rep2.sig), + (b"zeta\0", inp.gpr.z), + (b"trfra", inp.gpr.trainfrac), + (b"speci", " ".join(inp.system.species)), #str (and list str) + (b"nspe1", " ".join(inp.descriptor.rep1.neighspe)), + (b"nspe2", " ".join(inp.descriptor.rep2.neighspe)), + (b"dfbas", inp.qm.dfbasis) + ] + begin_of_block = SALTED_file.tell() + for key, value in inputs: + if debug: print(key, value) + write_key5(SALTED_file, key) + if isinstance(value, bool): + SALTED_file.write(i32(int(types_dict["bool"]))) + SALTED_file.write(sbool(int(value))) + elif isinstance(value, int): + SALTED_file.write(i32(int(types_dict["int32"]))) + SALTED_file.write(i32(value)) + elif isinstance(value, float): + SALTED_file.write(i32(int(types_dict["float64"]))) + SALTED_file.write(f64(value)) + elif isinstance(value, str): + SALTED_file.write(i32(int(types_dict["str"]))) + enc = value.encode("utf-8") + SALTED_file.write(i32(len(enc))) + SALTED_file.write(enc) + else: + print(key, value) + raise ValueError("Unknown type") + write_chunk_location(SALTED_file, "CONFG", begin_of_block) + + +def preprocess_shells(basis): + contractions_per_shell = [] + angular_momenta_per_shell = [] + coeffs_per_shell = [] + exponents_per_shell = [] + for shell in basis: + angular_momenta_per_shell.append(shell[0]) + contractions_per_shell.append(len(shell) - 1) + for exp,coef in shell[1:]: + exponents_per_shell.append(exp) + coeffs_per_shell.append(coef) + return (contractions_per_shell, angular_momenta_per_shell, + coeffs_per_shell, exponents_per_shell) + + + +#FORMAT FOR BASIS: +#Blocks start with the element symbol, then lines of the format: +# l, exponent, coefficient, shell number +#H +# 0 , 12.01 , 1.0000, 0 +# 0 , 145.01 , 0.5000, 1 +# 0 , 10.0 , 0.5000, 1 + +#He +# 0 , 12, 0.5, 0 +#... +#Helper function to read a basis, that is not supplied by PySCF, from a text file. +def read_new_basis(filename: str, symbol: str) -> list[list]: + with open(filename, "r", encoding="utf-8") as f: + content = f.read() + blocks = [block for block in content.split("\n\n")] + + for block in blocks: + lines = block.split("\n") + + atom_symbol = lines[0].strip() + + if atom_symbol != symbol: + continue + + basis = [] + last_shell = None + + for line in lines[1:]: + if line.startswith("#"): + continue + l_str, exp_str, coef_str, shell_str = [x.strip() for x in line.split(",")] + l = int(l_str) + exp = float(exp_str) + coef = float(coef_str) + shell = int(shell_str) + + if shell != last_shell: + basis.append([l]) + last_shell = shell + + basis[-1].append([exp, coef]) + + return basis + + raise ValueError(f"Symbol {symbol!r} not found in file") + + +#Format: +#TYPE_OF_DATA (4 bytes, int32) +#nElements (4 bytes, int32) +#For each element: + #Element ID (4 bytes, int32) + #FOR EACH ARRAY iN [CONTRACTIONS, ANGULAR MOMENTA, EXPONENTS, COEFFS]: + #N_DIMS (4 bytes, int32) + #DIMS (nDims*4 bytes, int32) + #DATA (dim1*8 bytes, float64) +def pack_basis(SALTED_file, inp, debug: bool = False): + try: + from pyscf import gto + from pyscf.data.elements import ELEMENTS + ELEMENTS_TO_NUM = {x: i for i,x in enumerate(ELEMENTS)} + except ImportError: + print("PySCF not found, cannot include basis sets in SALTED file, skipping basis packing") + return + + begin_of_block = SALTED_file.tell() + basis_name = inp.qm.dfbasis + symbols = inp.system.species + basis = {} + for symbol in symbols: + try: + basis[symbol] = gto.uncontract(gto.basis.load(basis_name, symb=symbol)) + except gto.basis.BasisNotFoundError: + try: + basis[symbol] = read_new_basis("additional_basis", symbol) + except ValueError as e: + print(f"Basis Set for symbol {symbol!r} not found in PySCF or additional_basis file, skipping basis packing") + return + SALTED_file.write(i32(int(types_dict["float64"]))) # Data type + SALTED_file.write(i32(int(len(symbols)))) # Number of elements (blocks to read after this) + for elem in symbols: + SALTED_file.write(i32(int(ELEMENTS_TO_NUM[elem]))) + shells = basis[elem] + (contractions_per_shell, + angular_momenta_per_shell, + coeffs_per_shell, + exponents_per_shell) = preprocess_shells(shells) + write_data_head(SALTED_file, np.array(contractions_per_shell)) + SALTED_file.write(np.array(contractions_per_shell, dtype=" Dict[int, np.ndarray]: + """ + Compute equivariant descriptors for all atoms. + + Supports both: + - Config mode: loads Wigner matrices from files + - Standalone mode: uses provided Wigner matrices + """ + nspe1 = len(neighspe1) + nspe2 = len(neighspe2) + ndata = len(conf_range) + natmax = max(natoms) if len(natoms) > 0 else 1 + natoms_total = sum(natoms) + + start_featomic = time.time() + + # Compute atomic representations + omega1 = sph_utils.get_representation_coeffs( + frames, rep1, HP1, rank, neighspe1, species, nang1, nrad1, natoms_total + ) + if sph_utils.reps_equivalent(rep1, neighspe1, HP1, rep2, neighspe2, HP2): + omega2 = omega1 + else: + omega2 = sph_utils.get_representation_coeffs( + frames, rep2, HP2, rank, neighspe2, species, nang2, nrad2, natoms_total + ) + + # Reshape arrays for Fortran indexing + v1 = np.transpose(omega1, (1, 3, 0, 2)).copy() + v2 = np.transpose(omega2, (1, 3, 0, 2)).copy() + + # Compute equivariant descriptors for each lambda + pvec = {} + for lam in range(lmax_max + 1): + if rank == 0: + print(f"lambda = {lam}", flush=True) + + llmax, llvec = sph_utils.get_angular_indexes_symmetric(lam, nang1, nang2) + + # Load Wigner-3J symbols + if wigners_dir is not None: + # Config mode: load from file + wigner3j = np.loadtxt(osp.join( + wigners_dir, f"wigner_lam-{lam}_lmax1-{nang1}_lmax2-{nang2}.dat" + )) + elif wigners_list is not None: + # Standalone mode: use provided list + if lam < len(wigners_list): + wigner3j = wigners_list[lam] + else: + if rank == 0: + print(f"Warning: Missing Wigner data for lambda={lam}", flush=True) + continue + else: + raise ValueError("Either wigners_dir or wigners_list must be provided") + + # Complex to real transformation + c2r = sph_utils.complex_to_real_transformation([2 * lam + 1])[0] + + # Compute combined features + if sparsify and vfps is not None and lam in vfps: + featsize = nspe1 * nspe2 * nrad1 * nrad2 * llmax + nfps_val = len(vfps[lam]) + p = sph_utils.equicombsparse_numba( + natoms_total, nang1, nang2, nspe1 * nrad1, nspe2 * nrad2, + v1, v2, wigner3j, llmax, llvec, lam, c2r, featsize, nfps_val, vfps[lam] + ) + featsize = ncut + else: + featsize = nspe1 * nspe2 * nrad1 * nrad2 * llmax + p = sph_utils.equicomb_numba( + natoms_total, nang1, nang2, nspe1 * nrad1, nspe2 * nrad2, + v1, v2, wigner3j, llmax, llvec, lam, c2r, featsize + ) + + # Store descriptors + if lam == 0: + p = p.reshape(natoms_total, featsize) + pvec[lam] = np.zeros((ndata, natmax, featsize)) + else: + p = p.reshape(natoms_total, 2 * lam + 1, featsize) + pvec[lam] = np.zeros((ndata, natmax, 2 * lam + 1, featsize)) + + j = 0 + for i, iconf in enumerate(conf_range): + for iat in range(natoms[iconf]): + pvec[lam][i, iat] = p[j] + j += 1 + + if rank == 0: + print(f"featomic time (sec) = {time.time() - start_featomic}", flush=True) + + return pvec + + +def compute_equivariant_descriptors_response( + frames: List, + conf_range: List[int], + natoms: np.ndarray, + atomic_symbols: List[List[str]], + rep1: str, + rep2: str, + HP1: Dict, + HP2: Dict, + nang1: int, + nang2: int, + nrad1: int, + nrad2: int, + neighspe1: List[str], + neighspe2: List[str], + species: List[str], + lmax: Dict, + lmax_max: int, + wigners_dir: Optional[str] = None, + wigners_list: Optional[List[np.ndarray]] = None, + wigners_antisymm_dir: Optional[str] = None, + wigners_antisymm_list: Optional[List[np.ndarray]] = None, + rank: int = 0) -> Tuple[Dict[int, np.ndarray], Dict[int, np.ndarray]]: + """ + Compute equivariant descriptors for density-response prediction. + + Returns both symmetric and antisymmetric descriptors. + """ + nspe1 = len(neighspe1) + nspe2 = len(neighspe2) + ndata = len(conf_range) + natmax = max(natoms) if len(natoms) > 0 else 1 + natoms_total = sum(natoms) + + start_featomic = time.time() + + # Compute atomic representations + omega1 = sph_utils.get_representation_coeffs( + frames, rep1, HP1, rank, neighspe1, species, nang1, nrad1, natoms_total + ) + if sph_utils.reps_equivalent(rep1, neighspe1, HP1, rep2, neighspe2, HP2): + omega2 = omega1 + else: + omega2 = sph_utils.get_representation_coeffs( + frames, rep2, HP2, rank, neighspe2, species, nang2, nrad2, natoms_total + ) + + # Reshape arrays for Fortran indexing + v1 = np.transpose(omega1, (1, 3, 0, 2)).copy() + v2 = np.transpose(omega2, (1, 3, 0, 2)).copy() + + lmax_max_response = lmax_max + 1 + + # Compute symmetric equivariant features + power = {} + for lam in range(lmax_max_response + 1): + if rank == 0: + print(f"lambda (symmetric) = {lam}", flush=True) + + llmax, llvec = sph_utils.get_angular_indexes_symmetric(lam, nang1, nang2) + + # Load Wigner-3J symbols + if wigners_dir is not None: + wigner3j = np.loadtxt(os.path.join( + wigners_dir, f"wigner_lam-{lam}_lmax1-{nang1}_lmax2-{nang2}.dat" + )) + elif wigners_list is not None: + if lam < len(wigners_list): + wigner3j = wigners_list[lam] + else: + if rank == 0: + print(f"Warning: Missing Wigner data for lambda={lam}", flush=True) + continue + else: + raise ValueError("Either wigners_dir or wigners_list must be provided") + + c2r = sph_utils.complex_to_real_transformation([2 * lam + 1])[0] + + # Compute symmetric features + featsize = nspe1 * nspe2 * nrad1 * nrad2 * llmax + p = equicombnonorm( + natoms_total, nang1, nang2, nspe1 * nrad1, nspe2 * nrad2, + v1, v2, wigner3j, llmax, llvec, lam, c2r, featsize + ) + + # Store descriptors + if lam == 0: + p = p.reshape(natoms_total, featsize) + power[lam] = np.zeros((ndata, natmax, featsize)) + else: + p = p.reshape(natoms_total, 2 * lam + 1, featsize) + power[lam] = np.zeros((ndata, natmax, 2 * lam + 1, featsize)) + + j = 0 + for i, iconf in enumerate(conf_range): + for iat in range(natoms[iconf]): + power[lam][i, iat] = p[j] + j += 1 + + # Compute antisymmetric equivariant features + power_antisymm = {} + for lam in range(1, lmax_max_response): + if rank == 0: + print(f"lambda (antisymmetric) = {lam}", flush=True) + + llmax, llvec = sph_utils.get_angular_indexes_antisymmetric(lam, nang1, nang2) -def build(): + # Load Wigner-3J symbols + if wigners_antisymm_dir is not None: + wigner3j = np.loadtxt(os.path.join( + wigners_antisymm_dir, f"wigner_antisymm_lam-{lam}_lmax1-{nang1}_lmax2-{nang2}.dat" + )) + elif wigners_antisymm_list is not None and len(wigners_antisymm_list) > 0: + if lam < len(wigners_antisymm_list): + wigner3j = wigners_antisymm_list[lam] + else: + if rank == 0: + print(f"Warning: Missing antisymmetric Wigner data for lambda={lam}", flush=True) + continue + else: + if rank == 0: + print(f"Warning: Antisymmetric Wigner data not provided, skipping lambda={lam}", flush=True) + continue + + wigdim = wigner3j.size + c2r = sph_utils.complex_to_real_transformation([2 * lam + 1])[0] + + # Compute antisymmetric features + featsize = nspe1 * nspe2 * nrad1 * nrad2 * llmax + p = antiequicombnonorm( + natoms_total, nang1, nang2, nspe1 * nrad1, nspe2 * nrad2, + v1, v2, wigner3j, llmax, llvec, lam, c2r, featsize + ) + p = p.reshape(natoms_total, 2 * lam + 1, featsize) + power_antisymm[lam] = np.zeros((ndata, natmax, 2 * lam + 1, featsize)) + + j = 0 + for i, iconf in enumerate(conf_range): + for iat in range(natoms[iconf]): + power_antisymm[lam][i, iat] = p[j] + j += 1 + + if rank == 0: + print(f"featomic time (sec) = {time.time() - start_featomic}", flush=True) + + return power, power_antisymm + + +def compute_prediction( + iconf_idx: int, + atomic_symbols: List[List[str]], + natoms: np.ndarray, + lmax: Dict, + nmax: Dict, + species: List[str], + psi_nm: Dict, + weights: np.ndarray, + cart: Optional[str] = None, + average: bool = False, + av_coefs: Optional[Dict] = None): + # Compute size + Tsize = 0 + for iat in range(natoms[iconf_idx]): + spe = atomic_symbols[iconf_idx][iat] + for l in range(lmax[spe] + 1): + for n in range(nmax[(spe, l)]): + Tsize += 2 * l + 1 + # Compute predictions per channel + C = {} + ispe = {} + isize = 0 + for spe in species: + ispe[spe] = 0 + for l in range(lmax[spe] + 1): + for n in range(nmax[(spe, l)]): + if cart is not None: + Mcut = psi_nm[(cart, spe, l)].shape[1] + C[(spe, l, n)] = np.dot(psi_nm[(cart, spe, l)], weights[isize:isize + Mcut]) + else: + Mcut = psi_nm[(spe, l)].shape[1] + C[(spe, l, n)] = np.dot(psi_nm[(spe, l)], weights[isize:isize + Mcut]) + isize += Mcut + + # Fill prediction vector + pred_coefs = np.zeros(Tsize) + Av_coeffs = np.zeros(Tsize) if average else None + + i = 0 + for iat in range(natoms[iconf_idx]): + spe = atomic_symbols[iconf_idx][iat] + for l in range(lmax[spe] + 1): + for n in range(nmax[(spe, l)]): + pred_coefs[i:i + 2 * l + 1] = C[(spe, l, n)][ + ispe[spe] * (2 * l + 1):ispe[spe] * (2 * l + 1) + 2 * l + 1 + ] + if average and l == 0 and av_coefs is not None: + Av_coeffs[i] = av_coefs[spe][n] + i += 2 * l + 1 + ispe[spe] += 1 + + # Add averages + if average and Av_coeffs is not None: + pred_coefs += Av_coeffs + + return pred_coefs + +def compute_density_descriptor_structure( + iconf: int, + i_local: Optional[int], + conf_range: List[int], + atom_idx: Dict, + natom_dict: Dict, + lmax: Dict, + species: List[str], + zeta: float, + pvec: Dict[int, np.ndarray], + power_env_sparse: Dict, + Vmat: Dict, + Mspe: Dict, + average: bool = False, + av_coefs: Optional[Dict] = None) -> np.ndarray: + """ + Compute predictions for a single structure. + """ + if i_local is None: + iconf_idx = conf_range.index(iconf) if isinstance(conf_range, list) else iconf + else: + iconf_idx = i_local + + # Compute kernels and projections + psi_nm = {} + for spe in species: + atom_indices = atom_idx[(iconf, spe)] + # Lambda = 0 + if (0, spe) in power_env_sparse: + if zeta == 1: + psi_nm[(spe, 0)] = np.dot( + pvec[0][iconf_idx, atom_indices], + power_env_sparse[(0, spe)].T + ) + else: + kernel0_nm = np.dot( + pvec[0][iconf_idx, atom_indices], + power_env_sparse[(0, spe)].T + ) + kernel_nm = kernel0_nm ** zeta + psi_nm[(spe, 0)] = np.dot(kernel_nm, Vmat[(0, spe)]) + + # Lambda > 0 + for lam in range(1, lmax[spe] + 1): + if (lam, spe) not in power_env_sparse: + continue + + featsize = pvec[lam].shape[-1] + + if zeta == 1: + psi_nm[(spe, lam)] = np.dot( + pvec[lam][iconf_idx, atom_indices].reshape( + natom_dict[(iconf, spe)] * (2 * lam + 1), featsize + ), + power_env_sparse[(lam, spe)].T + ) + else: + kernel_nm = np.dot( + pvec[lam][iconf_idx, atom_indices].reshape( + natom_dict[(iconf, spe)] * (2 * lam + 1), featsize + ), + power_env_sparse[(lam, spe)].T + ) + kernel_nm_blocks = kernel_nm.reshape( + natom_dict[(iconf, spe)], 2 * lam + 1, Mspe[spe], 2 * lam + 1 + ) + kernel_nm_blocks *= kernel0_nm[:, np.newaxis, :, np.newaxis] ** (zeta - 1) + kernel_nm = kernel_nm_blocks.reshape( + natom_dict[(iconf, spe)] * (2 * lam + 1), Mspe[spe] * (2 * lam + 1) + ) + psi_nm[(spe, lam)] = np.dot(kernel_nm, Vmat[(lam, spe)]) + + return psi_nm + +def compute_density_response_descriptor_structure( + iconf: int, + i_local: Optional[int], + conf_range: List[int], + atom_idx: Dict, + natom_dict: Dict, + lmax: Dict, + species: List[str], + zeta: float, + power: Dict[int, np.ndarray], + power_antisymm: Dict[int, np.ndarray], + power_env_sparse: Dict, + power_env_sparse_antisymm: Dict, + Vmat: Dict, + Mspe: Dict, + cart: List[str], + saltedpath: str, + saltedname: str, + Menv: int, + verbose: bool = False, + alpha_only: bool = False, + qmcode: str = "cp2k" +): + + if i_local is None: + iconf_idx = conf_range.index(iconf) if isinstance(conf_range, list) else iconf + else: + iconf_idx = i_local + + psi_nm = {} + psi_nm_cart = {} + # Compute kernels and RKHS descriptors + for ic in cart: + for spe in species: + for lam in range(lmax[spe] + 1): + psi_nm_cart[(ic, spe, lam)] = np.zeros( + (natom_dict[(iconf, spe)] * (2 * lam + 1), Vmat[(lam, spe)].shape[-1]) + ) + + for spe in species: + + start_kernel_0 = time.time() + + Mcut = {} + Mcutsize = {} + for lam in range(lmax[spe] + 1): + frac = np.exp(-0.05 * lam**2) + Mcut[lam] = int(round(Mspe[spe] * frac)) + Mcutsize[lam] = Mcut[lam] * 3 * (2 * lam + 1) + + # lam=0 + kernel0_nm = np.dot( + power[0][iconf_idx, atom_idx[(iconf, spe)]], power_env_sparse[(0, spe)].T + ) + kernel_nm = np.dot( + power[1][iconf_idx, atom_idx[(iconf, spe)]].reshape(natom_dict[(iconf, spe)] * 3, power[1].shape[-1]), + power_env_sparse[(1, spe)].T, + ) + + kernel_nm_blocks = kernel_nm.reshape(natom_dict[(iconf, spe)], 3, Mspe[spe], 3) + kernel_nm_blocks *= kernel0_nm[:, np.newaxis, :, np.newaxis] ** (zeta - 1) + kernel_nm = kernel_nm_blocks.reshape(natom_dict[(iconf, spe)] * 3, Mspe[spe] * 3) + kernel_nm = kernel_nm[:, :Mcutsize[0]] + + kernel0_nn_diag = np.sum(power[0][iconf_idx, atom_idx[(iconf, spe)]] ** 2, axis=1) + kernel_nn_diag = ( + power[1][iconf_idx, atom_idx[(iconf, spe)]] + @ power[1][iconf_idx, atom_idx[(iconf, spe)]].transpose(0, 2, 1) + ) + kernel_nn_diag = kernel_nn_diag * kernel0_nn_diag[:, np.newaxis, np.newaxis] ** (zeta - 1) + normfact = np.sqrt(np.sum(kernel_nn_diag**2, axis=(1, 2))) + + normfact_sparse = np.load( + os.path.join(saltedpath, f"normfacts_{saltedname}", f"M{Menv}_zeta{zeta}", f"normfact_spe-{spe}_lam-{0}.npy") + ) + knorm = kernelnorm(natom_dict[(iconf, spe)], Mcut[0], 3, normfact, normfact_sparse, np.real(kernel_nm)) + kernel_nm = knorm + + psi_nm[(spe, 0)] = np.real(np.dot(kernel_nm, Vmat[(0, spe)])) + + psi_nm_reshaped = psi_nm[(spe, 0)].reshape(natom_dict[(iconf, spe)], 3, psi_nm[(spe, 0)].shape[-1]) + for ik in range(3): + psi_nm_cart[(cart[ik], spe, 0)][:natom_dict[(iconf, spe)]] = psi_nm_reshaped[:, ik] + + if verbose: + print("kernel lam=0 time (sec) = ", time.time() - start_kernel_0, flush=True) + start_kernel_lam = time.time() + + if alpha_only and qmcode == "cp2k": + lmax[spe] = 1 + + # lam>0 + for lam in range(1,lmax[spe]+1): + + Msize = Mspe[spe] * 3 * (2 * lam + 1) + Nsize = natom_dict[(iconf, spe)] * 3 * (2 * lam + 1) + kernel_nm = np.zeros((Nsize,Msize),complex) + kernel_nn_diag = np.zeros((Nsize,3*(2*lam+1)),complex) + + # Perform CG combination + for L in [lam-1,lam,lam+1]: + + #print("L=", L) + + c2r = sph_utils.complex_to_real_transformation([2*L+1])[0] + + # compute complex descriptor for the given L + if L == lam: + pimag = power_antisymm[L][iconf_idx, atom_idx[(iconf, spe)]] + featsize = pimag.shape[-1] + pimag = pimag.reshape(natom_dict[(iconf, spe)], 2 * L + 1, featsize) + pimag = np.transpose(pimag, (1, 0, 2)).reshape(2 * L + 1, natom_dict[(iconf, spe)] * featsize) + preal = np.zeros_like(pimag) + else: + preal = power[L][iconf_idx, atom_idx[(iconf, spe)]] + featsize = preal.shape[-1] + preal = preal.reshape(natom_dict[(iconf, spe)], 2 * L + 1, featsize) + preal = np.transpose(preal, (1, 0, 2)).reshape(2 * L + 1, natom_dict[(iconf, spe)] * featsize) + pimag = np.zeros_like(preal) + + ptemp = preal + 1j * pimag + pcmplx = np.dot(np.conj(c2r.T),ptemp).reshape(2*L+1,natom_dict[(iconf_idx,spe)],featsize) + pcmplx = np.transpose(pcmplx,(1,0,2)).reshape(natom_dict[(iconf_idx,spe)]*(2*L+1),featsize) + + # compute complex sparse descriptor for the given L + if L == lam: + pimag = power_env_sparse_antisymm[(L, spe)] + featsize = pimag.shape[-1] + pimag = pimag.reshape(Mspe[spe], 2 * L + 1, featsize) + pimag = np.transpose(pimag, (1, 0, 2)).reshape(2 * L + 1, Mspe[spe] * featsize) + preal = np.zeros_like(pimag) + else: + preal = power_env_sparse[(L, spe)] + featsize = preal.shape[-1] + preal = preal.reshape(Mspe[spe], 2 * L + 1, featsize) + preal = np.transpose(preal, (1, 0, 2)).reshape(2 * L + 1, Mspe[spe] * featsize) + pimag = np.zeros_like(preal) + + ptemp = preal + 1j * pimag + pcmplx_sparse = np.dot(np.conj(c2r.T),ptemp).reshape(2*L+1,Mspe[spe],featsize) + pcmplx_sparse = np.transpose(pcmplx_sparse,(1,0,2)).reshape(Mspe[spe]*(2*L+1),featsize) + + # compute complex K_nm kernel + knm = np.dot(pcmplx, np.conj(pcmplx_sparse).T) + + # load the relevant CG coefficients + cgcoefs = np.loadtxt(os.path.join(saltedpath, "wigners", f"cg_response_lam-{lam}_L-{L}.dat")) + + k0 = kernel0_nm**(zeta-1) + cgkernel = kernelequicomb(natom_dict[(iconf_idx,spe)],Mspe[spe],lam,1,L,Nsize,Msize,len(cgcoefs),cgcoefs,knm,k0) + kernel_nm += cgkernel + + # compute complex K_nn kernel + pcmplx = pcmplx.reshape(natom_dict[(iconf_idx,spe)],2*L+1,featsize) + knn_diag = pcmplx @ np.conj(pcmplx).transpose(0,2,1) + knn_diag = knn_diag.reshape(natom_dict[(iconf_idx,spe)]*(2*L+1),2*L+1) + k0 = kernel0_nn_diag**(zeta-1) + cgkernel = kernelequicomb(natom_dict[(iconf_idx,spe)],1,lam,1,L,Nsize,3*(2*lam+1),len(cgcoefs),cgcoefs,knn_diag,k0[:,np.newaxis]) + kernel_nn_diag += cgkernel + + kernel_nm = kernel_nm[:, :Mcutsize[lam]] + + # compute complex to real transformation matrix for lam X 1 tensor product space + A = sph_utils.complex_to_real_transformation([2 * lam + 1])[0] + B = sph_utils.complex_to_real_transformation([3])[0] + c2r = np.zeros((3*(2*lam+1),3*(2*lam+1)),complex) + j1 = 0 + for i1 in range(2*lam+1): + j2 = 0 + for i2 in range(2*lam+1): + c2r[j1:j1+3,j2:j2+3] = A[i1,i2] * B + j2 += 3 + j1 += 3 + + # make k_NM real + ktemp1 = np.dot( + c2r, + np.transpose( + kernel_nm.reshape(natom_dict[(iconf, spe)], 3 * (2 * lam + 1), Mcutsize[lam]), + (1, 0, 2), + ).reshape(3 * (2 * lam + 1), natom_dict[(iconf, spe)] * Mcutsize[lam]), + ) + ktemp2 = np.transpose( + ktemp1.reshape(3 * (2 * lam + 1), natom_dict[(iconf, spe)], Mcutsize[lam]), (1, 0, 2) + ).reshape(Nsize, Mcutsize[lam]) + kernel_nm = np.dot( + ktemp2.reshape(Nsize, Mcut[lam], 3 * (2 * lam + 1)).reshape(Nsize, Mcut[lam], 3 * (2 * lam + 1)), + np.conj(c2r).T, + ).reshape(Nsize, Mcut[lam], 3 * (2 * lam + 1)).reshape(Nsize, Mcutsize[lam]) + + + # make k_NN_diag real and compute normalization factor + ktemp1 = np.dot( + c2r, + np.transpose( + kernel_nn_diag.reshape(natom_dict[(iconf, spe)], 3 * (2 * lam + 1), 3 * (2 * lam + 1)), + (1, 0, 2), + ).reshape(3 * (2 * lam + 1), natom_dict[(iconf, spe)] * 3 * (2 * lam + 1)), + ) + ktemp2 = np.transpose( + ktemp1.reshape(3 * (2 * lam + 1), natom_dict[(iconf, spe)], 3 * (2 * lam + 1)), (1, 0, 2) + ).reshape(Nsize, 3 * (2 * lam + 1)) + kernel_nn_diag = np.real(np.dot(ktemp2, np.conj(c2r).T)).reshape( + natom_dict[(iconf, spe)], 3 * (2 * lam + 1), 3 * (2 * lam + 1) + ) + normfact = np.sqrt(np.sum(kernel_nn_diag**2, axis=(1, 2))) + + normfact_sparse = np.load( + os.path.join(saltedpath, f"normfacts_{saltedname}", f"M{Menv}_zeta{zeta}", f"normfact_spe-{spe}_lam-{lam}.npy") + ) + knorm = kernelnorm( + natom_dict[(iconf, spe)], Mcut[lam], 3 * (2 * lam + 1), normfact, normfact_sparse, np.real(kernel_nm) + ) + kernel_nm = knorm + + # project kernel on the RKHS + psi_nm[(spe, lam)] = np.real(np.dot(kernel_nm, Vmat[(lam, spe)])) + + psi_nm_reshaped = psi_nm[(spe, lam)].reshape( + natom_dict[(iconf, spe)] * (2 * lam + 1), 3, psi_nm[(spe, lam)].shape[-1] + ) + for ik in range(3): + psi_nm_cart[(cart[ik], spe, lam)][: natom_dict[(iconf, spe)] * (2 * lam + 1)] = psi_nm_reshaped[:, ik] + + if inp.salted.verbose: + print("kernel lam>0 time (sec) = ",time.time()-start_kernel_lam,flush=True) + + return psi_nm_cart + + +def save_pred_descriptor( + data: Dict[int, np.ndarray], + config_range: List[int], + natoms: List[int], + dpath: str +): + """Save the descriptor data of the prediction dataset.""" + assert len(config_range) == len(natoms), ( + f"Length mismatch: {len(config_range)} vs {len(natoms)}" + ) + + for lam, data_this_lam in data.items(): + assert data_this_lam.shape[0] == len(config_range), ( + f"First dimension mismatch at lambda={lam}" + ) + + for idx, idx_in_full_dataset in enumerate(config_range): + this_data: Dict[str, np.ndarray] = {} + this_natoms = natoms[idx] + for lam, data_this_lam in data.items(): + this_data[f"lam{lam}"] = data_this_lam[idx, :this_natoms] + + with open(osp.join(dpath, f"descriptor_{idx_in_full_dataset + 1}.npz"), "wb") as f: + np.savez(f, **this_data) + + +# ============================================================================ +# CONFIG MODE (Full SALTED workflow) +# ============================================================================ +def predict_config_mode(): + """ + Full SALTED prediction workflow using config file. + + This maintains all features of the original prediction.py: + - Config file parsing + - MPI parallelization + - Density and density-response predictions + - CP2K integration for properties + - Directory structure loading + """ inp = ParseConfig().parse_input() (saltedname, saltedpath, saltedtype, filename, species, average, @@ -41,16 +717,16 @@ def build(): if filename_pred == PLACEHOLDER or predname == PLACEHOLDER: raise ValueError( - "No prediction file and name provided, " - "please specify the entry named `prediction.filename` and `prediction.predname` in the input file." + "No prediction file and name provided. " + "Specify 'prediction.filename' and 'prediction.predname' in input file." ) comm, size, rank, parallel = detect_mpi() - species, lmax, nmax, lmax_max, nnmax, ndata, atomic_symbols, natoms, natmax = read_system(filename_pred, species, dfbasis) - atom_idx, natom_dict = get_atom_idx(ndata,natoms,species,atomic_symbols) - - bohr2angs = 0.529177210670 + species, lmax, nmax, lmax_max, nnmax, ndata, atomic_symbols, natoms, natmax = read_system( + filename_pred, species, dfbasis + ) + atom_idx, natom_dict = get_atom_idx(ndata, natoms, species, atomic_symbols) if rank == 0: print(f"The dataset contains {ndata} frames.") @@ -59,17 +735,16 @@ def build(): if parallel: check_MPI_tasks_count(comm, ndata, "predicting structures") conf_range = distribute_jobs(comm, list(range(ndata))) - ndata = len(conf_range) # update ndata for each mpi task + ndata = len(conf_range) natmax = max(natoms[conf_range]) if inp.salted.verbose: - print(f"Task {rank} handles the following structures: {format_index_ranges(conf_range,True)}", flush=True) + print(f"Task {rank} handles: {format_index_ranges(conf_range, True)}", flush=True) else: conf_range = list(range(ndata)) - natoms_total = sum(natoms[conf_range]) - - reg_log10_intstr = str(int(np.log10(regul))) # for consistency + + reg_log10_intstr = str(int(np.log10(regul))) - # load regression weights + # Load regression weights ntrain = int(Ntrain * trainfrac) weights = np.load(osp.join( saltedpath, @@ -78,543 +753,478 @@ def build(): f"weights_N{ntrain}_reg{reg_log10_intstr}.npy" )) - if qmcode=="cp2k": - # Initialize calculation of density/density-response moments - charge_integrals,dipole_integrals = init_moments(inp,species,lmax,nmax,rank) + if qmcode == "cp2k": + charge_integrals, dipole_integrals = init_moments(inp, species, lmax, nmax, rank) - # base directory path for this prediction + # Setup output directory pdir = osp.join( saltedpath, f"predictions_{saltedname}_{predname}" ) - dirpath = osp.join(pdir, + dirpath = osp.join( + pdir, f"M{Menv}_zeta{zeta}", f"N{ntrain}_reg{reg_log10_intstr}", ) - # Create directory for predictions if rank == 0: if not os.path.exists(dirpath): os.makedirs(dirpath, exist_ok=True) - if saltedtype=="density-response": - for icart in ["x","y","z"]: - cartpath = os.path.join(dirpath, f"{icart}") + if saltedtype == "density-response": + for icart in ["x", "y", "z"]: + cartpath = os.path.join(dirpath, icart) if not os.path.exists(cartpath): os.mkdir(cartpath) + if parallel: comm.Barrier() - # Initialize files for derived properties - if qmcode=="cp2k": - if saltedtype=="density": - qfile = init_property_file("charges",saltedpath,pdir,Menv,zeta,ntrain,reg_log10_intstr,rank,size,comm) - dfile = init_property_file("dipoles",saltedpath,pdir,Menv,zeta,ntrain,reg_log10_intstr,rank,size,comm) - if saltedtype=="density-response": - pfile = init_property_file("polarizabilities",saltedpath,pdir,Menv,zeta,ntrain,reg_log10_intstr,rank,size,comm) + # Initialize property files for CP2K + qfile, dfile, pfile = None, None, None + if qmcode == "cp2k" : + if saltedtype == "density": + qfile = init_property_file("charges", saltedpath, pdir, Menv, zeta, ntrain, reg_log10_intstr, rank, size, comm) + dfile = init_property_file("dipoles", saltedpath, pdir, Menv, zeta, ntrain, reg_log10_intstr, rank, size, comm) + if saltedtype == "density-response": + pfile = init_property_file("polarizabilities", saltedpath, pdir, Menv, zeta, ntrain, reg_log10_intstr, rank, size, comm) start = time.time() - start_featomic = time.time() - # Read frames - frames = read(filename_pred,":") + frames = read(filename_pred, ":") frames = [frames[i] for i in conf_range] - # Compute atom-density spherical expansion coefficients - omega1 = sph_utils.get_representation_coeffs( - frames, rep1, HP1, rank, neighspe1, species, nang1, nrad1, natoms_total) - if sph_utils.reps_equivalent(rep1, neighspe1, HP1, rep2, neighspe2, HP2): - omega2 = omega1 - else: - omega2 = sph_utils.get_representation_coeffs( - frames, rep2, HP2, rank, neighspe2, species, nang2, nrad2, natoms_total) - - # Reshape arrays of expansion coefficients for optimal Fortran indexing - v1 = np.transpose(omega1,(1,3,0,2)).copy() - v2 = np.transpose(omega2,(1,3,0,2)).copy() - - print("featomic time (sec) = ",time.time()-start_featomic,flush=True) - - if saltedtype=="density": - - # Load feature space sparsification information if required + if saltedtype == "density": + # Load FPS information if required if sparsify: vfps = {} - for lam in range(lmax_max+1): + for lam in range(lmax_max + 1): vfps[lam] = np.load(osp.join( saltedpath, f"equirepr_{saltedname}", f"fps{ncut}-{lam}.npy" )) + else: + vfps = None - # Load training feature vectors and RKHS projection matrix - Vmat,Mspe,power_env_sparse = get_feats_projs(species,lmax) + # Load training features and RKHS projection matrix + Vmat, Mspe, power_env_sparse = get_feats_projs(species, lmax) # Load spherical averages if required + av_coefs_dict = {} if average: - av_coefs = {} for spe in species: - av_coefs[spe] = np.load(os.path.join(saltedpath, "coefficients", "averages", f"averages_{spe}.npy")) - - # Compute equivariant descriptors for each lambda value entering the SPH expansion of the electron density - pvec = {} - for lam in range(lmax_max+1): - - if rank == 0: print(f"lambda = {lam}") - - llmax, llvec = sph_utils.get_angular_indexes_symmetric(lam,nang1,nang2) - - # Load the relevant Wigner-3J symbols associated with the given triplet (lam, lmax1, lmax2) - wigner3j = np.loadtxt(osp.join( - saltedpath, "wigners", f"wigner_lam-{lam}_lmax1-{nang1}_lmax2-{nang2}.dat" - )) - wigdim = wigner3j.size - - # Compute complex to real transformation matrix for the given lambda value - c2r = sph_utils.complex_to_real_transformation([2*lam+1])[0] - - if sparsify: - - featsize = nspe1*nspe2*nrad1*nrad2*llmax - nfps = len(vfps[lam]) - p = sph_utils.equicombsparse_numba(natoms_total,nang1,nang2,nspe1*nrad1,nspe2*nrad2,v1,v2,wigner3j,llmax,llvec,lam,c2r,featsize,nfps,vfps[lam]) - featsize = ncut - - else: - - featsize = nspe1*nspe2*nrad1*nrad2*llmax - p = sph_utils.equicomb_numba(natoms_total,nang1,nang2,nspe1*nrad1,nspe2*nrad2,v1,v2,wigner3j,llmax,llvec,lam,c2r,featsize) - - # Fill vector of equivariant descriptor - if lam==0: - p = p.reshape(natoms_total,featsize) - pvec[lam] = np.zeros((ndata,natmax,featsize)) - else: - p = p.reshape(natoms_total,2*lam+1,featsize) - pvec[lam] = np.zeros((ndata,natmax,2*lam+1,featsize)) - - j = 0 - for i,iconf in enumerate(conf_range): - for iat in range(natoms[iconf]): - pvec[lam][i,iat] = p[j] - j += 1 - - """ save descriptor of the prediction dataset """ + av_coefs_dict[spe] = np.load( + os.path.join(saltedpath, "coefficients", "averages", f"averages_{spe}.npy") + ) + + # Compute equivariant descriptors + pvec = compute_equivariant_descriptors( + frames, conf_range, natoms, atomic_symbols, + rep1, rep2, HP1, HP2, + nang1, nang2, nrad1, nrad2, + neighspe1, neighspe2, species, lmax, lmax_max, + wigners_dir=osp.join(saltedpath, "wigners"), + sparsify=sparsify, + ncut=ncut, + vfps=vfps, + rank=rank, + ) + # Save descriptor if requested if inp.prediction.save_descriptor: if rank == 0: - print(f"Saving descriptor of the prediction dataset to dir {dirpath}") + print(f"Saving descriptor to {dirpath}", flush=True) save_pred_descriptor(pvec, conf_range, list(natoms[conf_range]), dirpath) - psi_nm = {} - for i,iconf in enumerate(conf_range): - - Tsize = 0 - for iat in range(natoms[iconf]): - spe = atomic_symbols[iconf][iat] - for l in range(lmax[spe]+1): - for n in range(nmax[(spe,l)]): - Tsize += 2*l+1 - - for spe in species: - - # lam = 0 - if zeta==1: - psi_nm[(spe,0)] = np.dot(pvec[0][i,atom_idx[(iconf,spe)]],power_env_sparse[(0,spe)].T) - else: - kernel0_nm = np.dot(pvec[0][i,atom_idx[(iconf,spe)]],power_env_sparse[(0,spe)].T) - kernel_nm = kernel0_nm**zeta - psi_nm[(spe,0)] = np.dot(kernel_nm,Vmat[(0,spe)]) - - # lam > 0 - for lam in range(1,lmax[spe]+1): - - featsize = pvec[lam].shape[-1] - if zeta==1: - psi_nm[(spe,lam)] = np.dot(pvec[lam][i,atom_idx[(iconf,spe)]].reshape(natom_dict[(iconf,spe)]*(2*lam+1),featsize),power_env_sparse[(lam,spe)].T) - else: - kernel_nm = np.dot(pvec[lam][i,atom_idx[(iconf,spe)]].reshape(natom_dict[(iconf,spe)]*(2*lam+1),featsize),power_env_sparse[(lam,spe)].T) - kernel_nm_blocks = kernel_nm.reshape(natom_dict[(iconf,spe)], 2*lam+1, Mspe[spe], 2*lam+1) - kernel_nm_blocks *= kernel0_nm[:, np.newaxis, :, np.newaxis] ** (zeta - 1) - kernel_nm = kernel_nm_blocks.reshape(natom_dict[(iconf,spe)]*(2*lam+1), Mspe[spe]*(2*lam+1)) - psi_nm[(spe,lam)] = np.dot(kernel_nm,Vmat[(lam,spe)]) - - # compute predictions per channel - C = {} - ispe = {} - isize = 0 - for spe in species: - ispe[spe] = 0 - for l in range(lmax[spe]+1): - for n in range(nmax[(spe,l)]): - Mcut = psi_nm[(spe,l)].shape[1] - C[(spe,l,n)] = np.dot(psi_nm[(spe,l)],weights[isize:isize+Mcut]) - isize += Mcut - - # init averages array if asked - if average: - Av_coeffs = np.zeros(Tsize) - - # fill vector of predictions - i = 0 - pred_coefs = np.zeros(Tsize) - for iat in range(natoms[iconf]): - spe = atomic_symbols[iconf][iat] - for l in range(lmax[spe]+1): - for n in range(nmax[(spe,l)]): - pred_coefs[i:i+2*l+1] = C[(spe,l,n)][ispe[spe]*(2*l+1):ispe[spe]*(2*l+1)+2*l+1] - if average and l==0: - Av_coeffs[i] = av_coefs[spe][n] - i += 2*l+1 - ispe[spe] += 1 - - # add back spherical averages if required - if average: - pred_coefs += Av_coeffs - - if qmcode=="cp2k": - # Compute charges and dipole moments - charge, dipole = compute_charge_and_dipole(frames[iconf],inp.qm.pseudocharge,natoms[iconf],atomic_symbols[iconf],lmax,nmax,species,charge_integrals,dipole_integrals,pred_coefs,average) - print(iconf+1,charge,file=qfile) - print(iconf+1,dipole["x"],dipole["y"],dipole["z"],file=dfile) + # Compute predictions + for i, iconf in enumerate(conf_range): + psi_nm = compute_density_descriptor_structure( + iconf, i, conf_range, atom_idx, natom_dict, + lmax, species, zeta, pvec, power_env_sparse, Vmat, Mspe, + average=average, av_coefs=av_coefs_dict + ) - # save predicted coefficients - np.savetxt(osp.join(dirpath, f"COEFFS-{iconf+1}.dat"), pred_coefs) - - - elif saltedtype=="density-response": - - start_load = time.time() - - # Load training feature vectors and RKHS projection matrix - Vmat,Mspe,power_env_sparse,power_env_sparse_antisymm = get_feats_projs_response(species,lmax) - - print("loading time (sec) = ",time.time()-start_load,flush=True) - + pred_coefs = compute_prediction( + iconf_idx=i, + atomic_symbols=atomic_symbols, + natoms=natoms, + lmax=lmax, + nmax=nmax, + species=species, + psi_nm=psi_nm, + weights=weights, + average=average, + av_coefs=av_coefs_dict) + + if qmcode == "cp2k": + charge, dipole = compute_charge_and_dipole( + frames[i], inp.qm.pseudocharge, natoms[iconf], atomic_symbols[iconf], + lmax, nmax, species, charge_integrals, dipole_integrals, pred_coefs, average + ) + print(iconf + 1, charge, file=qfile) + print(iconf + 1, dipole["x"], dipole["y"], dipole["z"], file=dfile) + + np.savetxt(osp.join(dirpath, f"COEFFS-{iconf + 1}.dat"), pred_coefs) + + elif saltedtype == "density-response": + # Load training features and RKHS projection matrix (response version) + Vmat, Mspe, power_env_sparse, power_env_sparse_antisymm = get_feats_projs_response(species, lmax) + lmax_max += 1 - cart = ["y","z","x"] - - start_feat = time.time() - - # Compute equivariant features for the given structure - power = {} - for lam in range(lmax_max+1): - - [llmax,llvec] = sph_utils.get_angular_indexes_symmetric(lam,nang1,nang2) - - # Load the relevant Wigner-3J symbols associated with the given triplet (lam, lmax1, lmax2) - wigner3j = np.loadtxt(os.path.join( - saltedpath, "wigners", f"wigner_lam-{lam}_lmax1-{nang1}_lmax2-{nang2}.dat" - )) - wigdim = wigner3j.size - - # Compute complex to real transformation matrix for the given lambda value - c2r = sph_utils.complex_to_real_transformation([2*lam+1])[0] - - # Perform symmetry-adapted combination following Eq.S19 of Grisafi et al., PRL 120, 036002 (2018) - featsize = nspe1*nspe2*nrad1*nrad2*llmax - p = equicombnonorm(natoms_total,nang1,nang2,nspe1*nrad1,nspe2*nrad2,v1,v2,wigner3j,llmax,llvec,lam,c2r,featsize) + + # Compute symmetric and antisymmetric equivariant descriptors + power, power_antisymm = compute_equivariant_descriptors_response( + frames, conf_range, natoms, atomic_symbols, + rep1, rep2, HP1, HP2, + nang1, nang2, nrad1, nrad2, + neighspe1, neighspe2, species, lmax, lmax_max, + wigners_dir=osp.join(saltedpath, "wigners"), + wigners_antisymm_dir=osp.join(saltedpath, "wigners"), + rank=rank, + ) - # Fill vector of equivariant descriptor - if lam==0: - p = p.reshape(natoms_total,featsize) - power[lam] = np.zeros((ndata,natmax,featsize)) - else: - p = p.reshape(natoms_total,2*lam+1,featsize) - power[lam] = np.zeros((ndata,natmax,2*lam+1,featsize)) + for i, iconf in enumerate(conf_range): + if rank == 0 or i % 10 == 0: + print(f"Predicting structure {iconf + 1}/{len(conf_range)}...", flush=True) + + start = time.time() + psi_nm_cart = compute_density_response_descriptor_structure( + iconf, i, conf_range, atom_idx, natom_dict, + lmax, species, zeta, + power, power_antisymm, power_env_sparse, power_env_sparse_antisymm, + Vmat, Mspe, cart, saltedpath, saltedname, Menv, + verbose=inp.salted.verbose, + alpha_only=alpha_only, + qmcode=qmcode, + ) + + for icart in ["x", "y", "z"]: + pred_coefs = compute_prediction( + i, atomic_symbols, natoms, lmax, nmax, species, psi_nm_cart, weights, + cart=icart, + average=average, + av_coefs=av_coefs + ) + np.savetxt(os.path.join(dirpath, icart, f"COEFFS-{iconf + 1}.dat"), pred_coefs) + + if rank == 0 or i % 10 == 0: + print(f"done in {time.time() - start:.2f} seconds.", flush=True) - j = 0 - for i,iconf in enumerate(conf_range): - for iat in range(natoms[iconf]): - power[lam][i,iat] = p[j] - j += 1 - # Compute antisymmetric equivariant features for the given structure - power_antisymm = {} - for lam in range(1,lmax_max): + if qmcode=="cp2k": + # Compute polarizability + alpha = compute_polarizability(frames[iconf],natoms[iconf],atomic_symbols[iconf],lmax,nmax,species,charge_integrals,dipole_integrals,pred_coefs) - [llmax,llvec] = sph_utils.get_angular_indexes_antisymmetric(lam,nang1,nang2) + # Save polarizabilities + print(iconf+1, alpha[("x","x")], alpha[("x","y")], alpha[("x","z")], + alpha[("y","x")], alpha[("y","y")], alpha[("y","z")], + alpha[("z","x")], alpha[("z","y")], alpha[("z","z")], + file=pfile) - # Load the relevant Wigner-3J symbols associated with the given triplet (lam, lmax1, lmax2) - wigner3j = np.loadtxt(os.path.join( - saltedpath, "wigners", f"wigner_antisymm_lam-{lam}_lmax1-{nang1}_lmax2-{nang2}.dat" - )) - wigdim = wigner3j.size + if inp.salted.verbose: + print("prediction time (sec) = ", time.time() - start, flush=True) - # Compute complex to real transformation matrix for the given lambda value - c2r = sph_utils.complex_to_real_transformation([2*lam+1])[0] + # Close property files + if qmcode == "cp2k": + if qfile is not None: + qfile.close() + if dfile is not None: + dfile.close() + if pfile is not None: + pfile.close() - # Perform symmetry-adapted combination following Eq.S19 of Grisafi et al., PRL 120, 036002 (2018) - featsize = nspe1*nspe2*nrad1*nrad2*llmax - p = antiequicombnonorm(natoms_total,nang1,nang2,nspe1*nrad1,nspe2*nrad2,v1,v2,wigner3j,llmax,llvec,lam,c2r,featsize) + if rank == 0: + print(f"\nTotal time: {(time.time() - start):.2f} s") - # Fill vector of equivariant descriptor - p = p.reshape(natoms_total,2*lam+1,featsize) - power_antisymm[lam] = np.zeros((ndata,natmax,2*lam+1,featsize)) - j = 0 - for i,iconf in enumerate(conf_range): - for iat in range(natoms[iconf]): - power_antisymm[lam][i,iat] = p[j] - j += 1 - - print("features time (sec) =", time.time()-start_feat,flush=True) +# ============================================================================ +# STANDALONE MODE (Binary model file) +# ============================================================================ +def predict_standalone_mode(model_file: str, xyz_file: str, output_dir: str = None): + """ + Standalone prediction using .salted model file with MPI support. + """ + # Detect MPI + comm, size, rank, parallel = detect_mpi() - psi_nm = {} - psi_nm_cart = {} - for i,iconf in enumerate(conf_range): + if rank == 0: + print(f"Loading model from {model_file}...") + + model = read_model.read_salted_model(model_file) + + # Extract configuration + config = model['config'] + species_str = config['speci'] + species = species_str.split() + + # Extract hyperparameters + nang1 = int(config['nang1']) + nang2 = int(config['nang2']) + nrad1 = int(config['nrad1']) + nrad2 = int(config['nrad2']) + rcut1 = float(config['rcut1']) + rcut2 = float(config['rcut2']) + sig1 = float(config['sig1']) + sig2 = float(config['sig2']) + zeta = int(config['zeta']) + use_average = bool(config['averg']) + use_sparsify = bool(config['spars']) + ncut = int(config.get('ncut', 0)) + dfbasis = config['dfbas'] + saltedtype = config.get('predtype', 'density') # Default to density if not specified + + neighspe1_str = config.get('nspe1', species_str) + neighspe2_str = config.get('nspe2', species_str) + neighspe1 = neighspe1_str.split() + neighspe2 = neighspe2_str.split() + + # Representation types + rep1 = "rho" + rep2 = "rho" - Tsize = 0 - for iat in range(natoms[iconf]): - spe = atomic_symbols[iconf][iat] - for l in range(lmax[spe]+1): - for n in range(nmax[(spe,l)]): - Tsize += 2*l+1 + if rank == 0: + print(f"Model configuration:") + print(f" Species: {species}") + print(f" Basis: {dfbasis}") + print(f" Prediction type: {saltedtype}") + print(f" nang1={nang1}, nang2={nang2}, nrad1={nrad1}, nrad2={nrad2}") + print(f" zeta={zeta}, sparsify={use_sparsify}, average={use_average}") + + # Read system + if rank == 0: + print(f"\nReading structures from {xyz_file}...") + + species, lmax, nmax, lmax_max, nnmax, ndata, atomic_symbols, natoms, natmax = read_system( + xyz_file, species, dfbasis, basis_data=model.get('basis') + ) + + if rank == 0: + print(f" Found {ndata} structures") + print(f" Max atoms per structure: {natmax}") - # Compute kernels and RKHS descriptors - for ic in cart: - for spe in species: - for lam in range(lmax[spe]+1): - psi_nm_cart[(ic,spe,lam)] = np.zeros((natom_dict[(iconf,spe)]*(2*lam+1),Vmat[(lam,spe)].shape[-1])) + # Setup output directory + if output_dir is None: + output_dir = "predictions" + + if rank == 0: + os.makedirs(output_dir, exist_ok=True) + print(f" Output directory: {output_dir}") + + if parallel: + comm.Barrier() - for spe in species: + # Distribute structures across MPI ranks + if parallel: + check_MPI_tasks_count(comm, ndata, "predicting structures") + conf_range = distribute_jobs(comm, list(range(ndata))) + if rank == 0: + print(f"Distributed structures across {size} MPI tasks") + print(f"Task {rank} handles: {format_index_ranges(conf_range, True)}", flush=True) + else: + conf_range = list(range(ndata)) - start_kernel_0 = time.time() + # Get atom indices for all structures + atom_idx, natom_dict = get_atom_idx(ndata, natoms, species, atomic_symbols) - Mcut = {} - Mcutsize = {} - for lam in range(lmax[spe]+1): - frac = np.exp(-0.05*lam**2) - Mcut[lam] = int(round(Mspe[spe]*frac)) - Mcutsize[lam] = Mcut[lam]*3*(2*lam+1) + # Read frames + frames = read(xyz_file, ":") + frames_local = [frames[i] for i in conf_range] + + # Extract model data + weights = model['weights'] + wigners = model.get('wigners', []) + averages = model.get('averages', {}) + fps_data = model.get('fps', []) + feats = model['feats'] + projectors = model['projectors'] + + # Setup hyperparameters + HP1 = { + "cutoff": {"radius": rcut1, "smoothing": {"type": "ShiftedCosine", "width": 0.1}}, + "density": {"type": "Gaussian", "width": sig1}, + "basis": { + "type": "TensorProduct", + "max_angular": nang1, + "radial": {"type": "Gto", "max_radial": nrad1 - 1}, + "spline_accuracy": 1e-06 + } + } + + HP2 = { + "cutoff": {"radius": rcut2, "smoothing": {"type": "ShiftedCosine", "width": 0.1}}, + "density": {"type": "Gaussian", "width": sig2}, + "basis": { + "type": "TensorProduct", + "max_angular": nang2, + "radial": {"type": "Gto", "max_radial": nrad2 - 1}, + "spline_accuracy": 1e-06 + } + } + + # Prepare FPS indices if using sparsification + vfps = {} + if use_sparsify: + for lam in range(lmax_max + 1): + if lam < len(fps_data): + vfps[lam] = fps_data[lam] - # lam=0 - kernel0_nm = np.dot(power[0][i,atom_idx[(iconf,spe)]],power_env_sparse[(0,spe)].T) - kernel_nm = np.dot(power[1][i,atom_idx[(iconf,spe)]].reshape(natom_dict[(iconf,spe)]*3,power[1].shape[-1]),power_env_sparse[(1,spe)].T) + if rank == 0: + print("\nComputing atomic representations...") + + start_time = time.time() + + if saltedtype == "density": + # Compute equivariant descriptors for density prediction + pvec = compute_equivariant_descriptors( + frames_local, conf_range, natoms, atomic_symbols, + rep1, rep2, HP1, HP2, + nang1, nang2, nrad1, nrad2, + neighspe1, neighspe2, species, lmax, lmax_max, + lcut=None, + wigners_list=wigners, + sparsify=use_sparsify, + ncut=ncut, + vfps=vfps if use_sparsify else None, + rank=rank, + ) - kernel_nm_blocks = kernel_nm.reshape(natom_dict[(iconf,spe)], 3, Mspe[spe], 3) - kernel_nm_blocks *= kernel0_nm[:, np.newaxis, :, np.newaxis] ** (zeta - 1) - kernel_nm = kernel_nm_blocks.reshape(natom_dict[(iconf,spe)] * 3, Mspe[spe] * 3) - kernel_nm = kernel_nm[:,:Mcutsize[0]] - - kernel0_nn_diag = np.sum(power[0][i,atom_idx[(iconf,spe)]]**2,axis=1) - kernel_nn_diag = power[1][i,atom_idx[(iconf,spe)]] @ power[1][i,atom_idx[(iconf,spe)]].transpose(0,2,1) - kernel_nn_diag = kernel_nn_diag * kernel0_nn_diag[:,np.newaxis,np.newaxis]**(zeta-1) - normfact = np.sqrt(np.sum(kernel_nn_diag**2,axis=(1,2))) - - normfact_sparse = np.load(os.path.join(saltedpath, f"normfacts_{saltedname}", f"M{Menv}_zeta{zeta}", f"normfact_spe-{spe}_lam-{0}.npy")) - knorm = kernelnorm(natom_dict[(iconf,spe)],Mcut[0],3,normfact,normfact_sparse,np.real(kernel_nm)) - kernel_nm = knorm - - psi_nm[(spe,0)] = np.real(np.dot(kernel_nm,Vmat[(0,spe)])) - - psi_nm_reshaped = psi_nm[(spe, 0)].reshape(natom_dict[(iconf,spe)], 3, psi_nm[(spe, 0)].shape[-1]) - for ik in range(3): - psi_nm_cart[(cart[ik], spe, 0)][:natom_dict[(iconf,spe)]] = psi_nm_reshaped[:, ik] - - if inp.salted.verbose: - print("kernel lam=0 time (sec) = ",time.time()-start_kernel_0,flush=True) - start_kernel_lam = time.time() - - if alpha_only and qmcode=="cp2k": - lmax[spe] = 1 - - # lam>0 - for lam in range(1,lmax[spe]+1): - - Msize = Mspe[spe]*3*(2*lam+1) - Nsize = natom_dict[(iconf,spe)]*3*(2*lam+1) - kernel_nm = np.zeros((Nsize,Msize),complex) - kernel_nn_diag = np.zeros((Nsize,3*(2*lam+1)),complex) - - # Perform CG combination - for L in [lam-1,lam,lam+1]: - - #print("L=", L) - - c2r = sph_utils.complex_to_real_transformation([2*L+1])[0] - - # compute complex descriptor for the given L - if L==lam: - pimag = power_antisymm[L][i,atom_idx[(iconf,spe)]] - featsize = pimag.shape[-1] - pimag = pimag.reshape(natom_dict[(iconf,spe)],2*L+1,featsize) - pimag = np.transpose(pimag,(1,0,2)).reshape(2*L+1,natom_dict[(iconf,spe)]*featsize) - preal = np.zeros_like(pimag) - else: - preal = power[L][i,atom_idx[(iconf,spe)]] - featsize = preal.shape[-1] - preal = preal.reshape(natom_dict[(iconf,spe)],2*L+1,featsize) - preal = np.transpose(preal,(1,0,2)).reshape(2*L+1,natom_dict[(iconf,spe)]*featsize) - pimag = np.zeros_like(preal) - - ptemp = preal + 1j * pimag - pcmplx = np.dot(np.conj(c2r.T),ptemp).reshape(2*L+1,natom_dict[(iconf,spe)],featsize) - pcmplx = np.transpose(pcmplx,(1,0,2)).reshape(natom_dict[(iconf,spe)]*(2*L+1),featsize) - - # compute complex sparse descriptor for the given L - if L==lam: - pimag = power_env_sparse_antisymm[(L,spe)] - featsize = pimag.shape[-1] - pimag = pimag.reshape(Mspe[spe],2*L+1,featsize) - pimag = np.transpose(pimag,(1,0,2)).reshape(2*L+1,Mspe[spe]*featsize) - preal = np.zeros_like(pimag) - else: - preal = power_env_sparse[(L,spe)] - featsize = preal.shape[-1] - preal = preal.reshape(Mspe[spe],2*L+1,featsize) - preal = np.transpose(preal,(1,0,2)).reshape(2*L+1,Mspe[spe]*featsize) - pimag = np.zeros_like(preal) - - ptemp = preal + 1j * pimag - pcmplx_sparse = np.dot(np.conj(c2r.T),ptemp).reshape(2*L+1,Mspe[spe],featsize) - pcmplx_sparse = np.transpose(pcmplx_sparse,(1,0,2)).reshape(Mspe[spe]*(2*L+1),featsize) - - # compute complex K_nm kernel - knm = np.dot(pcmplx,np.conj(pcmplx_sparse).T) - - # load the relevant CG coefficients - cgcoefs = np.loadtxt(os.path.join(saltedpath, "wigners", f"cg_response_lam-{lam}_L-{L}.dat")) - - k0 = kernel0_nm**(zeta-1) - cgkernel = kernelequicomb(natom_dict[(iconf,spe)],Mspe[spe],lam,1,L,Nsize,Msize,len(cgcoefs),cgcoefs,knm,k0) - kernel_nm += cgkernel - - # compute complex K_nn kernel - pcmplx = pcmplx.reshape(natom_dict[(iconf,spe)],2*L+1,featsize) - knn_diag = pcmplx @ np.conj(pcmplx).transpose(0,2,1) - knn_diag = knn_diag.reshape(natom_dict[(iconf,spe)]*(2*L+1),2*L+1) - k0 = kernel0_nn_diag**(zeta-1) - cgkernel = kernelequicomb(natom_dict[(iconf,spe)],1,lam,1,L,Nsize,3*(2*lam+1),len(cgcoefs),cgcoefs,knn_diag,k0[:,np.newaxis]) - kernel_nn_diag += cgkernel - - kernel_nm = kernel_nm[:,:Mcutsize[lam]] - - # compute complex to real transformation matrix for lam X 1 tensor product space - A = sph_utils.complex_to_real_transformation([2*lam+1])[0] - B = sph_utils.complex_to_real_transformation([3])[0] - c2r = np.zeros((3*(2*lam+1),3*(2*lam+1)),complex) - j1 = 0 - for i1 in range(2*lam+1): - j2 = 0 - for i2 in range(2*lam+1): - c2r[j1:j1+3,j2:j2+3] = A[i1,i2] * B - j2 += 3 - j1 += 3 - - # make k_NM real - ktemp1 = np.dot(c2r,np.transpose(kernel_nm.reshape(natom_dict[(iconf,spe)],3*(2*lam+1),Mcutsize[lam]),(1,0,2)).reshape(3*(2*lam+1),natom_dict[(iconf,spe)]*Mcutsize[lam])) - ktemp2 = np.transpose(ktemp1.reshape(3*(2*lam+1),natom_dict[(iconf,spe)],Mcutsize[lam]),(1,0,2)).reshape(Nsize,Mcutsize[lam]) - kernel_nm = np.dot(ktemp2.reshape(Nsize,Mcut[lam],3*(2*lam+1)).reshape(Nsize*Mcut[lam],3*(2*lam+1)),np.conj(c2r).T).reshape(Nsize,Mcut[lam],3*(2*lam+1)).reshape(Nsize,Mcutsize[lam]) - - - # make k_NN_diag real and compute normalization factor - ktemp1 = np.dot(c2r,np.transpose(kernel_nn_diag.reshape(natom_dict[(iconf,spe)],3*(2*lam+1),3*(2*lam+1)),(1,0,2)).reshape(3*(2*lam+1),natom_dict[(iconf,spe)]*3*(2*lam+1))) - ktemp2 = np.transpose(ktemp1.reshape(3*(2*lam+1),natom_dict[(iconf,spe)],3*(2*lam+1)),(1,0,2)).reshape(Nsize,3*(2*lam+1)) - kernel_nn_diag = np.real(np.dot(ktemp2,np.conj(c2r).T)).reshape(natom_dict[(iconf,spe)],3*(2*lam+1),3*(2*lam+1)) - normfact = np.sqrt(np.sum(kernel_nn_diag**2,axis=(1,2))) - - normfact_sparse = np.load(os.path.join(saltedpath, f"normfacts_{saltedname}", f"M{Menv}_zeta{zeta}", f"normfact_spe-{spe}_lam-{lam}.npy")) - knorm = kernelnorm(natom_dict[(iconf,spe)],Mcut[lam],3*(2*lam+1),normfact,normfact_sparse,np.real(kernel_nm)) - kernel_nm = knorm - - # project kernel on the RKHS - psi_nm[(spe,lam)] = np.real(np.dot(kernel_nm,Vmat[(lam,spe)])) - - psi_nm_reshaped = psi_nm[(spe, lam)].reshape(natom_dict[(iconf,spe)]*(2*lam+1), 3, psi_nm[(spe, lam)].shape[-1]) - for ik in range(3): - psi_nm_cart[(cart[ik], spe, lam)][:natom_dict[(iconf,spe)]*(2*lam+1)] = psi_nm_reshaped[:, ik] - - if inp.salted.verbose: - print("kernel lam>0 time (sec) = ",time.time()-start_kernel_lam,flush=True) - - start_pred = time.time() - - pred_coefs = {} - for icart in ["x","y","z"]: - - # compute predictions per channel - C = {} - ispe = {} - isize = 0 - for spe in species: - ispe[spe] = 0 - for l in range(lmax[spe]+1): - for n in range(nmax[(spe,l)]): - Mcut = psi_nm_cart[(icart,spe,l)].shape[1] - C[(spe,l,n)] = np.dot(psi_nm_cart[(icart,spe,l)],weights[isize:isize+Mcut]) - isize += Mcut - - # fill vector of predictions - i = 0 - pred_coefs[icart] = np.zeros(Tsize) - for iat in range(natoms[iconf]): - spe = atomic_symbols[iconf][iat] - for l in range(lmax[spe]+1): - for n in range(nmax[(spe,l)]): - pred_coefs[icart][i:i+2*l+1] = C[(spe,l,n)][ispe[spe]*(2*l+1):ispe[spe]*(2*l+1)+2*l+1] - i += 2*l+1 - ispe[spe] += 1 - - # save predicted coefficients - np.savetxt(osp.join(dirpath, f"{icart}", f"COEFFS-{iconf+1}.dat"), pred_coefs[icart]) + if rank == 0: + print(f" Feature computation time: {time.time() - start_time:.2f} s") + + # Load projectors and features + Vmat = {} + power_env_sparse = {} + Mspe = {} + + for spe in species: + if spe not in feats or spe not in projectors: + if rank == 0: + print(f"Warning: Missing data for species {spe}") + continue + + for lam in range(lmax[spe] + 1): + lam_str = str(lam) + if lam_str in feats[spe]: + power_env_sparse[(lam, spe)] = feats[spe][lam_str] + if lam_str in projectors[spe]: + Vmat[(lam, spe)] = projectors[spe][lam_str] + + if lam == 0: + Mspe[spe] = power_env_sparse[(lam, spe)].shape[0] + + # Compute predictions for density + if rank == 0: + print(f"\nComputing predictions for {ndata} structures...") + + for i, iconf in enumerate(conf_range): + if rank == 0 or i % 10 == 0: + print(f" [Rank {rank}] Structure {iconf + 1}/{ndata}...", end=" ", flush=True) - if qmcode=="cp2k": - # Compute polarizability - alpha = compute_polarizability(frames[iconf],natoms[iconf],atomic_symbols[iconf],lmax,nmax,species,charge_integrals,dipole_integrals,pred_coefs) + start = time.time() + + psi_nm = compute_density_descriptor_structure( + iconf, i, conf_range, atom_idx, natom_dict, + lmax, species, zeta, pvec, power_env_sparse, Vmat, Mspe, + average=use_average, av_coefs=averages if use_average else None + ) + + pred_coefs = compute_prediction( + i, atomic_symbols, natoms, lmax, nmax, species, psi_nm, weights, + average=use_average, + av_coefs=averages + ) + # Save predictions + output_file = os.path.join(output_dir, f"COEFFS-{iconf + 1}.dat") + np.savetxt(output_file, pred_coefs) + + if rank == 0 or i % 10 == 0: + print(f"done ({time.time() - start:.2f} s)") + + elif saltedtype == "density-response": + # Compute equivariant descriptors for density-response prediction + # Note: For density-response, wigners_antisymm_list would need to be available in the model + # Currently using empty list as placeholder + power, power_antisymm = compute_equivariant_descriptors_response( + frames_local, conf_range, natoms, atomic_symbols, + rep1, rep2, HP1, HP2, + nang1, nang2, nrad1, nrad2, + neighspe1, neighspe2, species, lmax, lmax_max, + wigners_list=wigners, + wigners_antisymm_list=[], # TODO: Extract from model when available + rank=rank, + ) - # Save polarizabilities - print(iconf+1, alpha[("x","x")], alpha[("x","y")], alpha[("x","z")], - alpha[("y","x")], alpha[("y","y")], alpha[("y","z")], - alpha[("z","x")], alpha[("z","y")], alpha[("z","z")], - file=pfile) + if rank == 0: + print(f" Feature computation time: {time.time() - start_time:.2f} s") + print("Density-response prediction not fully implemented in standalone mode yet") - if inp.salted.verbose: - print("prediction time (sec) = ",time.time()-start_pred,flush=True) + else: + raise ValueError(f"Unknown prediction type: {saltedtype}") - if qmcode == "cp2k": - if saltedtype=="density": - qfile.close() - dfile.close() - if saltedtype=="density-response": - pfile.close() + if parallel: + comm.Barrier() - if rank == 0: print(f"\ntotal time: {(time.time()-start):.2f} s") + if rank == 0: + print(f"\nPredictions saved to {output_dir}/") + + +# ============================================================================ +# CLI AND MAIN +# ============================================================================ +def main(): + parser = argparse.ArgumentParser( + description="Unified prediction script for SALTED models", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + + CONFIG MODE (full SALTED workflow): + python predict_unified.py + + STANDALONE MODE (binary model file): + python predict_unified.py --model model.salted --xyz structures.xyz + python predict_unified.py --model model.salted --xyz structures.xyz --output predictions/ + """ + ) + + parser.add_argument( + "--model", + type=str, + help="Path to .salted model file (standalone mode)" + ) + parser.add_argument( + "--xyz", + type=str, + help="Path to XYZ structures file (required for standalone mode)" + ) + parser.add_argument( + "--output", + type=str, + default="predictions", + help="Output directory for predictions (standalone mode, default: predictions)" + ) + args = parser.parse_args() -def save_pred_descriptor(data: dict[int, np.ndarray], config_range: list[int], natoms: list[int], dpath: str): - """Save the descriptor data of the prediction dataset. + if args.model and not args.xyz: + parser.error("--xyz is required when using --model") - Args: - data (dict[int, np.ndarray]): the descriptor data to be saved. - int -> lambda value, - np.ndarray -> descriptor data, shape (ndata, natmax, [2*lambda+1,] featsize) - natmax should be cut to the number of atoms in the structure (natoms[i]) - 2*lambda+1 is only for lambda > 0. - config_range (list[int]): the indices of the structures in the full dataset. - natoms (list[int]): the number of atoms in each structure. Should be the same length as config_range. - dpath (str): the directory to save the descriptor data. - Output: - The descriptor data of each structure is saved in a separate npz file in the directory dpath named as - "descriptor_{i}.npz", where i starts from 1. - Format: npz file with keys as lambda values and values as the descriptor data. - Values have shape (natom, [2*lambda+1,] featsize). 2*lambda+1 is only for lambda > 0. - """ - assert len(config_range) == len(natoms), f"The length of config_range and natoms should be the same, " \ - f"but get {config_range=} and {natoms=}." - for lam, data_this_lam in data.items(): - assert data_this_lam.shape[0] == len(config_range), \ - f"The first dimension of the descriptor data should be the same as the length of config_range, " \ - f"but at {lam=} get {data_this_lam.shape[0]=} and {len(config_range)=}." + if args.model: + # Run standalone mode + try: + predict_standalone_mode(args.model, args.xyz, args.output) + exit() + except Exception as e: + print(f"Error in standalone mode: {e}") + exit(1) - """ cut natmax to the number of atoms in the structure """ - for idx, idx_in_full_dataset in enumerate(config_range): - this_data: dict[int, np.ndarray] = dict() - this_natoms = natoms[idx] - for lam, data_this_lam in data.items(): - this_data[f"lam{lam}"] = data_this_lam[idx, :this_natoms] # shape (natom, [2*lambda+1,] featsize) - with open(osp.join(dpath, f"descriptor_{idx_in_full_dataset+1}.npz"), "wb") as f: # index starts from 1 - np.savez(f, **this_data) + # As fallback, also allow running config mode if no model file is provided + try: + predict_config_mode() + exit() + except Exception as e: + print(f"Error in config mode: {e}") if __name__ == "__main__": - build() + main() diff --git a/salted/pyscf/dm2df.py b/salted/pyscf/dm2df.py index fd3eed64..4cd1188a 100644 --- a/salted/pyscf/dm2df.py +++ b/salted/pyscf/dm2df.py @@ -62,13 +62,10 @@ def cal_df_coeffs( overlap_reordered[l1_indexes_to] = overlap[l1_indexes_from] overlap_reordered[:, l1_indexes_to] = overlap[:, l1_indexes_from] - # Compute density projections on auxiliary functions - proj_reordered = np.dot(overlap_reordered, coef_reordered) reorder_time = time.time() - reorder_time return { "coef": coef_reordered, - "proj": proj_reordered, "over": overlap_reordered, "pyscf_time": pyscf_time, "reorder_time": reorder_time, @@ -98,7 +95,7 @@ def main(geom_indexes: list[int] | None, num_threads: int = None): """check if all subdirectories exist, if not create them""" sub_dirs = [ osp.join(inp.salted.saltedpath, d) - for d in ("overlaps", "coefficients", "projections") + for d in ("overlaps", "coefficients") ] for sub_dir in sub_dirs: if not osp.exists(sub_dir): @@ -137,7 +134,6 @@ def main(geom_indexes: list[int] | None, num_threads: int = None): dm = np.load(osp.join(inp.qm.path2qm, "density_matrices", f"dm_conf{geom_idx+1}.npy")) reordered_data = cal_df_coeffs(atoms, inp.qm.qmbasis, ribasis, dm, irreps) np.save(osp.join(inp.salted.saltedpath, "coefficients", f"coefficients_conf{geom_idx}.npy"), reordered_data["coef"]) - np.save(osp.join(inp.salted.saltedpath, "projections", f"projections_conf{geom_idx}.npy"), reordered_data["proj"]) np.save(osp.join(inp.salted.saltedpath, "overlaps", f"overlap_conf{geom_idx}.npy"), reordered_data["over"]) pyscf_time += reordered_data["pyscf_time"] reorder_time += reordered_data["reorder_time"] diff --git a/salted/read_model.py b/salted/read_model.py new file mode 100644 index 00000000..656acc31 --- /dev/null +++ b/salted/read_model.py @@ -0,0 +1,410 @@ +import struct +import numpy as np +from typing import Dict, List, Tuple, Any, Optional + +# Unpacking functions for little-endian format +def read_u32(f): return struct.unpack(' str: + """Read a 5-byte key and strip null bytes""" + key_bytes = f.read(5) + return key_bytes.rstrip(b'\0').decode('utf-8') + +def read_data_head(f) -> Tuple[int, Tuple[int, ...]]: + """Read the dimensionality and shape of array data""" + ndims = read_i32(f) + dims = tuple(read_i32(f) for _ in range(ndims)) + return ndims, dims + +def read_header(f) -> Tuple[int, Dict[str, int]]: + """Read the file header and return version and block locations""" + # Read magic number + magic = f.read(5) + if magic != MAGIC_NUMBER: + raise ValueError(f"Invalid SALTED file: magic number is {magic!r}, expected {MAGIC_NUMBER!r}") + + # Read version + version = read_i32(f) + if version not in SUPPORTED_VERSIONS: + raise ValueError(f"Unsupported version {version}, supported versions: {SUPPORTED_VERSIONS}") + + # Read number of blocks + n_blocks = read_i32(f) + + # Read table of contents (block names and locations) + blocks = {} + for _ in range(n_blocks): + block_name = read_key5(f) + block_location = read_i32(f) + blocks[block_name] = block_location + + return version, blocks + +def read_averages(f) -> Dict[str, np.ndarray]: + """Read the AVERG block containing averages for each element""" + data_type = read_i32(f) + if types_dict[data_type] != "float64": + raise ValueError(f"Expected float64 data type for averages, got {types_dict[data_type]}") + + nfiles = read_i32(f) + averages = {} + + for _ in range(nfiles): + element = read_key5(f) + ndims, dims = read_data_head(f) + nbytes = np.prod(dims) * 8 + data = np.frombuffer(f.read(nbytes), dtype=' List[np.ndarray]: + """Read the WIG block containing Wigner matrices""" + data_type = read_i32(f) + if types_dict[data_type] != "float64": + raise ValueError(f"Expected float64 data type for wigners, got {types_dict[data_type]}") + + nfiles = read_i32(f) + wigners = [] + + for _ in range(nfiles): + ndims, dims = read_data_head(f) + nbytes = np.prod(dims) * 8 + data = np.frombuffer(f.read(nbytes), dtype=' List[np.ndarray]: + """Read the FPS block containing FPS indices""" + data_type = read_i32(f) + if types_dict[data_type] != "int64": + raise ValueError(f"Expected int64 data type for FPS, got {types_dict[data_type]}") + + nfiles = read_i32(f) + fps_data = [] + + for _ in range(nfiles): + ndims, dims = read_data_head(f) + nbytes = np.prod(dims) * 8 + data = np.frombuffer(f.read(nbytes), dtype=' Dict[str, Dict[str, np.ndarray]]: + """Read the PROJE block containing projectors""" + data_type = read_i32(f) + if types_dict[data_type] != "float64": + raise ValueError(f"Expected float64 data type for projectors, got {types_dict[data_type]}") + + nkeys = read_i32(f) + projectors = {} + + for _ in range(nkeys): + species = read_key5(f) + nlambda = read_i32(f) + projectors[species] = {} + + for _ in range(nlambda): + ndims, dims = read_data_head(f) + nbytes = np.prod(dims) * 8 + data = np.frombuffer(f.read(nbytes), dtype=' Dict[str, Dict[str, np.ndarray]]: + """Read the FEATS block containing sparse descriptors""" + data_type = read_i32(f) + if types_dict[data_type] != "float64": + raise ValueError(f"Expected float64 data type for FEATS, got {types_dict[data_type]}") + + nkeys = read_i32(f) + feats = {} + + for _ in range(nkeys): + species = read_key5(f) + nlambda = read_i32(f) + feats[species] = {} + + for _ in range(nlambda): + ndims, dims = read_data_head(f) + nbytes = np.prod(dims) * 8 + data = np.frombuffer(f.read(nbytes), dtype=' np.ndarray: + """Read the WEIGH block containing regression weights""" + data_type = read_i32(f) + if types_dict[data_type] != "float64": + raise ValueError(f"Expected float64 data type for weights, got {types_dict[data_type]}") + + nfiles = read_i32(f) + if nfiles != 1: + raise ValueError(f"Expected 1 weights file, got {nfiles}") + + ndims, dims = read_data_head(f) + nbytes = np.prod(dims) * 8 + weights = np.frombuffer(f.read(nbytes), dtype=' Dict[str, Any]: + """Read the CONFG block containing model configuration""" + config = {} + + # Read until we hit the end of the block or another block marker + # We need to read carefully since we don't know the block size + start_pos = f.tell() + + try: + while True: + # Try to read a key + pos_before = f.tell() + key = read_key5(f) + + # Check if this might be a block marker (next block) + # If we can't read a valid type, we've reached the end + try: + value_type = read_i32(f) + except struct.error: + f.seek(pos_before) + break + + if value_type not in types_dict: + # This is probably the start of the next block + f.seek(pos_before) + break + + type_name = types_dict[value_type] + + if type_name == "bool": + value = bool(read_bool(f)) + elif type_name == "int32": + value = read_i32(f) + elif type_name == "float64": + value = read_f64(f) + elif type_name == "str": + str_len = read_i32(f) + value = f.read(str_len).decode('utf-8') + else: + raise ValueError(f"Unexpected type in config: {type_name}") + + config[key] = value + except Exception as e: + # If we encounter an error, we might be at the end + pass + + return config + +def read_basis(f) -> Dict[str, Dict[str, np.ndarray]]: + """Read the BASIS block containing basis set information""" + try: + from pyscf.data.elements import ELEMENTS + except ImportError: + print("Warning: PySCF not found, element names won't be resolved") + ELEMENTS = None + + data_type = read_i32(f) + if types_dict[data_type] != "float64": + raise ValueError(f"Expected float64 data type for basis, got {types_dict[data_type]}") + + n_elements = read_i32(f) + basis = {} + + for _ in range(n_elements): + element_num = read_i32(f) + if ELEMENTS is not None and element_num < len(ELEMENTS): + element_symbol = ELEMENTS[element_num] + else: + element_symbol = f"Element_{element_num}" + + basis[element_symbol] = {} + + # Read contractions per shell + ndims, dims = read_data_head(f) + nbytes = np.prod(dims) * 4 + basis[element_symbol]['contractions'] = np.frombuffer(f.read(nbytes), dtype=' Dict[str, Any]: + """ + Read a SALTED model file and return all its contents. + + Parameters + ---------- + filename : str + Path to the .salted file + + Returns + ------- + model : dict + Dictionary containing all model data with keys: + - 'version': File format version + - 'config': Model configuration parameters + - 'averages': Average values per element + - 'wigners': Wigner matrices + - 'fps': FPS indices + - 'feats': Sparse descriptors + - 'projectors': Projector matrices + - 'weights': Regression weights + - 'basis': Basis set information (if available) + """ + model = {} + + with open(filename, 'rb') as f: + # Read header and get block locations + version, blocks = read_header(f) + model['version'] = version + + print(f"SALTED model version {version}") + print(f"Available blocks: {list(blocks.keys())}") + + # Read each block + if 'CONFG' in blocks: + print("Reading configuration...") + f.seek(blocks['CONFG']) + model['config'] = read_config(f) + + if 'AVERG' in blocks: + print("Reading averages...") + f.seek(blocks['AVERG']) + model['averages'] = read_averages(f) + + if 'WIG' in blocks: + print("Reading Wigner matrices...") + f.seek(blocks['WIG']) + model['wigners'] = read_wigners(f) + + if 'FPS' in blocks: + print("Reading FPS indices...") + f.seek(blocks['FPS']) + model['fps'] = read_fps(f) + + if 'FEATS' in blocks: + print("Reading sparse descriptors...") + f.seek(blocks['FEATS']) + model['feats'] = read_feats(f) + + if 'PROJ' in blocks: + print("Reading projectors...") + f.seek(blocks['PROJ']) + model['projectors'] = read_projectors(f) + + if 'WEIGH' in blocks: + print("Reading weights...") + f.seek(blocks['WEIGH']) + model['weights'] = read_weights(f) + + if 'BASIS' in blocks: + print("Reading basis sets...") + f.seek(blocks['BASIS']) + model['basis'] = read_basis(f) + + return model + +def print_model_summary(model: Dict[str, Any]): + """Print a summary of the loaded model""" + print("\n" + "="*60) + print("SALTED Model Summary") + print("="*60) + + print(f"\nVersion: {model.get('version', 'N/A')}") + + if 'config' in model: + print("\nConfiguration:") + for key, value in sorted(model['config'].items()): + print(f" {key}: {value}") + + if 'averages' in model: + print(f"\nAverages: {len(model['averages'])} elements") + for elem, data in model['averages'].items(): + print(f" {elem}: shape {data.shape}") + + if 'wigners' in model: + print(f"\nWigner matrices: {len(model['wigners'])} lambdas") + for i, w in enumerate(model['wigners']): + print(f" Lambda {i}: shape {w.shape}") + + if 'fps' in model: + print(f"\nFPS indices: {len(model['fps'])} arrays") + for i, fps in enumerate(model['fps']): + print(f" Array {i}: shape {fps.shape}") + + if 'feats' in model: + print(f"\nSparse descriptors: {len(model['feats'])} species") + for species, lambdas in model['feats'].items(): + print(f" {species}: {len(lambdas)} lambdas") + + if 'projectors' in model: + print(f"\nProjectors: {len(model['projectors'])} species") + for species, lambdas in model['projectors'].items(): + print(f" {species}: {len(lambdas)} lambdas") + + if 'weights' in model: + print(f"\nWeights: shape {model['weights'].shape}") + + if 'basis' in model: + print(f"\nBasis sets: {len(model['basis'])} elements") + for elem in model['basis'].keys(): + print(f" {elem}") + + print("="*60) + +if __name__ == "__main__": + import sys + + if len(sys.argv) < 2: + print("Usage: python read_model.py ") + sys.exit(1) + + filename = sys.argv[1] + print(f"Reading SALTED model from {filename}...") + + try: + model = read_salted_model(filename) + print_model_summary(model) + except Exception as e: + print(f"Error reading model: {e}") + import traceback + traceback.print_exc() + sys.exit(1) diff --git a/salted/sparsify_features.py b/salted/sparsify_features.py index b5868aaa..f2c17500 100644 --- a/salted/sparsify_features.py +++ b/salted/sparsify_features.py @@ -76,6 +76,7 @@ def build(): v1 = np.transpose(omega1,(1,3,0,2)).copy() v2 = np.transpose(omega2,(1,3,0,2)).copy() + del omega1, omega2 # Compute equivariant descriptors for each lambda value entering the SPH expansion of the electron density for lam in range(lmax_max+1): diff --git a/salted/sys_utils.py b/salted/sys_utils.py index 5c4bf3ab..0d0c79eb 100644 --- a/salted/sys_utils.py +++ b/salted/sys_utils.py @@ -54,13 +54,41 @@ def build_featomic_hyper_params(rep_cfg) -> dict: raise ValueError(f"Unknown representation type '{rep_cfg.type}': must be 'rho' or 'V'.") -def read_system(filename: str = None, spelist: list[str] = None, dfbasis: str = None): +def _basis_from_embedded_model(basis_data: dict[str, dict], spelist: list[str]): + """Convert embedded model basis data to the old (lmax, nmax) format.""" + lmax = {} + nmax = {} + + for spe in spelist: + if spe not in basis_data: + raise ValueError(f"Embedded basis data does not contain species {spe!r}") + + spe_data = basis_data[spe] + angular_momenta = np.asarray(spe_data["angular_momenta"], dtype=int).ravel() + if angular_momenta.size == 0: + raise ValueError(f"Embedded basis data for species {spe!r} is empty") + + spe_lmax = int(angular_momenta.max()) + lmax[spe] = spe_lmax + for l in range(spe_lmax + 1): + nmax[(spe, l)] = int(np.count_nonzero(angular_momenta == l)) + + return lmax, nmax + + +def read_system( + filename: str = None, + spelist: list[str] = None, + dfbasis: str = None, + basis_data: dict[str, dict] | None = None, +): """read a geometry file and return the formatted information Args: filename (str, optional): geometry file. Defaults to None. spelist (list[str], optional): list of species. Defaults to None. dfbasis (str, optional): density fitting basis. Defaults to None. + basis_data (dict, optional): embedded basis data loaded from a .salted model. Notes: By default (all parameters are None), it reads the geometry file for training dataset. @@ -78,9 +106,8 @@ def read_system(filename: str = None, spelist: list[str] = None, dfbasis: str = natmax (int): maximum number of atoms in the system """ - inp = ParseConfig().parse_input() - if (filename is None) and (spelist is None) and (dfbasis is None): + inp = ParseConfig().parse_input() filename = inp.system.filename spelist = inp.system.species dfbasis = inp.qm.dfbasis @@ -92,8 +119,11 @@ def read_system(filename: str = None, spelist: list[str] = None, dfbasis: str = "please check the docstring for more details." ) - # read basis - [lmax, nmax] = basis.basiset(dfbasis) + # read basis: prefer embedded model basis data when available, otherwise fall back to the named basis set + if basis_data is not None: + [lmax, nmax] = _basis_from_embedded_model(basis_data, spelist) + else: + [lmax, nmax] = basis.basiset(dfbasis) llist = [] nlist = [] for spe in spelist: @@ -203,6 +233,7 @@ def detect_mpi(): return comm, size, rank, parallel else: return None, 1, 0, False + def check_MPI_tasks_count(comm, num_items: int, item_name: str = "items"):