shnitsel-tools 0.0.1__tar.gz

shnitsel_tools-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 SHNITSEL
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

shnitsel_tools-0.0.1/PKG-INFO

@@ -0,0 +1,197 @@
+Metadata-Version: 2.4
+Name: shnitsel-tools
+Version: 0.0.1
+Summary: Display and interpret output of SHARC
+Author-email: Robin Curth <robin.curth@uni-leipzig.de>, Theodor Everley Röhrkasten <theodor.roehrkasten@fau.de>, Carolin Müller <carolin.cpc.mueller@fau.de>, Julia Westermayr <julia.westermayr@uni-leipzig.de>
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: scipy>=1.11.4
+Requires-Dist: matplotlib
+Requires-Dist: xarray
+Requires-Dist: rdkit
+Requires-Dist: ase
+Requires-Dist: numpy>=1.26.2
+Requires-Dist: pandas>=2.1.4
+Requires-Dist: h5netcdf>=1.3.0
+Requires-Dist: scikit-learn>=1.3.2
+Requires-Dist: jupyter>=1.0.0
+Requires-Dist: py3dmol>=2.0.4
+Requires-Dist: tqdm>=4.66.2
+Provides-Extra: extras
+Requires-Dist: plotly>=5.22.0; extra == "extras"
+Requires-Dist: seaborn>=0.13.1; extra == "extras"
+Requires-Dist: bokeh; extra == "extras"
+Provides-Extra: dev
+Requires-Dist: pytest>=5.2; extra == "dev"
+Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
+Requires-Dist: hypothesis; extra == "dev"
+Requires-Dist: tox>=4.10.0; extra == "dev"
+Requires-Dist: pre-commit>=2.20.0; extra == "dev"
+Requires-Dist: ruff==0.4.1; extra == "dev"
+Requires-Dist: mypy>=1.9.0; extra == "dev"
+Requires-Dist: pandas-stubs; extra == "dev"
+Requires-Dist: scipy-stubs; extra == "dev"
+Requires-Dist: types-tqdm; extra == "dev"
+Dynamic: license-file
+
+<div align="center">
+  <h1>shnitsel-tools</h1>
+  <img src="https://github.com/SHNITSEL/shnitsel-tools/blob/main/logo_shnitsel_tools.png" alt="SHNITSEL-TOOLS Logo" width="200px">
+  <h3>Surface Hopping Nested Instances Training Set for Excited-state Learning Tools</h3>
+  <br>
+  <a href="[https://zenodo.org/records/15482819](https://shnitsel.github.io/)">
+    <img src="https://img.shields.io/badge/Website-shnitsel.github.io-yellow.svg" alt="DOI">
+  </a>
+</div>
+
+--------------------
+
+## About
+
+`shnitsel-tools` is designed to to support the entire data lifecycle of surface hopping (SH) trajectory data upon simulation: data managment, storage, processing, visualization and interpretation. 
+The tool is compatible with surface hopping data generated using the software packages [SHARC 3/4](https://sharc-md.org/), [Newton-X](https://newtonx.org/), and [PyRAI2MD](https://github.com/lopez-lab/PyRAI2MD).
+The `shnitsel` code leverages Xarray to benefit from efficient multidimensional data handling, improved metadata management, and a structure that aligns naturally with the needs of quantum chemical datasets.
+
+## Usage
+
+The package is organized into ...
+
+> [!TIP]
+> ### Tutorials
+>
+> For a quick start, see the [tutorials](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials) directory,
+> which contains Jupyter Notebooks showing the workflow for parsing, writing and loading SHNITSEL databases as well as how to postprocess and visualize the respective data.
+> 
+> #### Collection & storage
+> - [parsing trajcetory and initial condition data obtained by SHARC](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials/0_1_sharc2hdf5.ipynb)
+> - [parsing trajectory data produced with Newton-X](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials/0_2_nx2hdf5.ipynb)
+> - [convert ASE databases](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials/0_4_ase2hdf5.ipynb)
+>
+> #### Management
+>
+> [exploration of electronic properties](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials/2_2_PS_explore.ipynb)
+>
+> #### Postprocessing & visualization of data
+> - [datasheet for trajectory data](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials/3_1_datasheet.ipynb)
+> - [principal component analysis and trajectory classification](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials/1_1_GS_PCA.ipynb)
+>
+> #### Walkthroughs
+>
+> (Link to folder, further info in README)
+>
+
+```bash
+shnitsel
+├── core
+│   ├── ase.py
+│   ├── datasheet
+│   │   ├── __init__.py
+│   │   ├── oop.py
+│   │   └── plot
+│   │       ├── colormaps.py
+│   │       ├── common.py
+│   │       ├── dip_trans_hist.py
+│   │       ├── hist.py
+│   │       ├── __init__.py
+│   │       ├── nacs_hist.py
+│   │       ├── per_state_hist.py
+│   │       ├── structure.py
+│   │       └── time.py
+│   ├── filter_unphysical.py
+│   ├── filtre.py
+│   ├── indexes.py
+│   ├── __init__.py
+│   ├── parse
+│   │   ├── common.py
+│   │   ├── __init__.py
+│   │   ├── nx.py
+│   │   ├── pyrai2md.py
+│   │   ├── sharc_icond.py
+│   │   ├── sharc_traj.py
+│   │   └── xyz.py
+│   ├── pca_biplot.py
+│   ├── plot
+│   │   ├── __init__.py
+│   │   ├── kde.py
+│   │   ├── p3mhelpers.py
+│   │   ├── select.py
+│   │   └── spectra3d.py
+│   ├── plotting.py
+│   ├── postprocess.py
+│   ├── spectra.py
+│   └── xrhelpers.py
+├── __init__.py
+├── plot
+│   └── __init__.py
+├── rd.py
+├── README.md
+└── xarray.py
+```
+
+## Installation
+
+You can create the environment with a custom path using one of the following methods:
+
+<details open>
+  <summary><strong>Option 1: Using `uv`</strong></summary>
+  We recommend to use the `uv` tool, available at https://docs.astral.sh/uv/.  
+  Run the following in the `shnitsel-tools` directory:
+
+  ```bash
+  uv venv  # create an environment under ./.venv
+  . .venv/bin/activate  # activate the new environment
+  uv pip install -e .  # install shnitsel in editable mode
+  ```
+
+  To install the optional development dependencies run
+
+  ```bash
+  uv pip install -e '.[dev]'
+  ```
+  
+</details>
+
+<details open>
+  <summary><strong>Option 2: Using the `--prefix` Flag</strong></summary>
+  
+  You can create the environment and specify the desired path by using the `conda env create` command with the `--prefix` flag:
+  
+  ```bash
+  conda env create --prefix /home/user/anaconda3/envs/shnitsel -f shnitsel-tools.yml
+  ```
+</details>
+
+<details>
+  <summary><strong>Option 3: Adding the Path to the .yml File</strong></summary>
+  
+  Alternatively, you can manually add the desired path to the shnitsel-tools.yml file and create the environment directly:
+    
+  1) Open the shnitsel-tools.yml file for editing:
+  
+  ```bash
+  vi shnitsel-tools.yml
+  ```
+  
+  2) Add the following line to the file:
+  
+  
+  ```
+  prefix: /home/user/anaconda3/envs/shnitsel
+  ```
+  
+  3) Create the environment with a custom path. 
+  
+  ```bash
+  conda env create -f shnitsel-rdkit.yml
+  ```
+</details>
+
+## Further Information
+
+[![Website](https://img.shields.io/badge/Website-shnitsel.github.io-yellow.svg)](https://shnitsel.github.io/)
+
+

shnitsel_tools-0.0.1/README.md

@@ -0,0 +1,157 @@
+<div align="center">
+  <h1>shnitsel-tools</h1>
+  <img src="https://github.com/SHNITSEL/shnitsel-tools/blob/main/logo_shnitsel_tools.png" alt="SHNITSEL-TOOLS Logo" width="200px">
+  <h3>Surface Hopping Nested Instances Training Set for Excited-state Learning Tools</h3>
+  <br>
+  <a href="[https://zenodo.org/records/15482819](https://shnitsel.github.io/)">
+    <img src="https://img.shields.io/badge/Website-shnitsel.github.io-yellow.svg" alt="DOI">
+  </a>
+</div>
+
+--------------------
+
+## About
+
+`shnitsel-tools` is designed to to support the entire data lifecycle of surface hopping (SH) trajectory data upon simulation: data managment, storage, processing, visualization and interpretation. 
+The tool is compatible with surface hopping data generated using the software packages [SHARC 3/4](https://sharc-md.org/), [Newton-X](https://newtonx.org/), and [PyRAI2MD](https://github.com/lopez-lab/PyRAI2MD).
+The `shnitsel` code leverages Xarray to benefit from efficient multidimensional data handling, improved metadata management, and a structure that aligns naturally with the needs of quantum chemical datasets.
+
+## Usage
+
+The package is organized into ...
+
+> [!TIP]
+> ### Tutorials
+>
+> For a quick start, see the [tutorials](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials) directory,
+> which contains Jupyter Notebooks showing the workflow for parsing, writing and loading SHNITSEL databases as well as how to postprocess and visualize the respective data.
+> 
+> #### Collection & storage
+> - [parsing trajcetory and initial condition data obtained by SHARC](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials/0_1_sharc2hdf5.ipynb)
+> - [parsing trajectory data produced with Newton-X](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials/0_2_nx2hdf5.ipynb)
+> - [convert ASE databases](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials/0_4_ase2hdf5.ipynb)
+>
+> #### Management
+>
+> [exploration of electronic properties](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials/2_2_PS_explore.ipynb)
+>
+> #### Postprocessing & visualization of data
+> - [datasheet for trajectory data](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials/3_1_datasheet.ipynb)
+> - [principal component analysis and trajectory classification](https://github.com/SHNITSEL/shnitsel-tools/blob/main/tutorials/1_1_GS_PCA.ipynb)
+>
+> #### Walkthroughs
+>
+> (Link to folder, further info in README)
+>
+
+```bash
+shnitsel
+├── core
+│   ├── ase.py
+│   ├── datasheet
+│   │   ├── __init__.py
+│   │   ├── oop.py
+│   │   └── plot
+│   │       ├── colormaps.py
+│   │       ├── common.py
+│   │       ├── dip_trans_hist.py
+│   │       ├── hist.py
+│   │       ├── __init__.py
+│   │       ├── nacs_hist.py
+│   │       ├── per_state_hist.py
+│   │       ├── structure.py
+│   │       └── time.py
+│   ├── filter_unphysical.py
+│   ├── filtre.py
+│   ├── indexes.py
+│   ├── __init__.py
+│   ├── parse
+│   │   ├── common.py
+│   │   ├── __init__.py
+│   │   ├── nx.py
+│   │   ├── pyrai2md.py
+│   │   ├── sharc_icond.py
+│   │   ├── sharc_traj.py
+│   │   └── xyz.py
+│   ├── pca_biplot.py
+│   ├── plot
+│   │   ├── __init__.py
+│   │   ├── kde.py
+│   │   ├── p3mhelpers.py
+│   │   ├── select.py
+│   │   └── spectra3d.py
+│   ├── plotting.py
+│   ├── postprocess.py
+│   ├── spectra.py
+│   └── xrhelpers.py
+├── __init__.py
+├── plot
+│   └── __init__.py
+├── rd.py
+├── README.md
+└── xarray.py
+```
+
+## Installation
+
+You can create the environment with a custom path using one of the following methods:
+
+<details open>
+  <summary><strong>Option 1: Using `uv`</strong></summary>
+  We recommend to use the `uv` tool, available at https://docs.astral.sh/uv/.  
+  Run the following in the `shnitsel-tools` directory:
+
+  ```bash
+  uv venv  # create an environment under ./.venv
+  . .venv/bin/activate  # activate the new environment
+  uv pip install -e .  # install shnitsel in editable mode
+  ```
+
+  To install the optional development dependencies run
+
+  ```bash
+  uv pip install -e '.[dev]'
+  ```
+  
+</details>
+
+<details open>
+  <summary><strong>Option 2: Using the `--prefix` Flag</strong></summary>
+  
+  You can create the environment and specify the desired path by using the `conda env create` command with the `--prefix` flag:
+  
+  ```bash
+  conda env create --prefix /home/user/anaconda3/envs/shnitsel -f shnitsel-tools.yml
+  ```
+</details>
+
+<details>
+  <summary><strong>Option 3: Adding the Path to the .yml File</strong></summary>
+  
+  Alternatively, you can manually add the desired path to the shnitsel-tools.yml file and create the environment directly:
+    
+  1) Open the shnitsel-tools.yml file for editing:
+  
+  ```bash
+  vi shnitsel-tools.yml
+  ```
+  
+  2) Add the following line to the file:
+  
+  
+  ```
+  prefix: /home/user/anaconda3/envs/shnitsel
+  ```
+  
+  3) Create the environment with a custom path. 
+  
+  ```bash
+  conda env create -f shnitsel-rdkit.yml
+  ```
+</details>
+
+## Further Information
+
+[![Website](https://img.shields.io/badge/Website-shnitsel.github.io-yellow.svg)](https://shnitsel.github.io/)
+
+

shnitsel_tools-0.0.1/pyproject.toml

@@ -0,0 +1,103 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "shnitsel-tools"
+version = "0.0.1"
+authors = [
+  { name="Robin Curth", email="robin.curth@uni-leipzig.de" },
+  { name="Theodor Everley Röhrkasten", email="theodor.roehrkasten@fau.de" },
+  { name="Carolin Müller", email="carolin.cpc.mueller@fau.de" },
+  { name="Julia Westermayr", email="julia.westermayr@uni-leipzig.de" },
+
+]
+description = "Display and interpret output of SHARC"
+readme = "README.md"
+requires-python = ">=3.10"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+  "scipy>=1.11.4",
+  "matplotlib",
+  "xarray",
+  "rdkit",
+  "ase",
+  "numpy>=1.26.2",
+  "pandas>=2.1.4",
+  "h5netcdf>=1.3.0",
+  "scikit-learn>=1.3.2",
+  "jupyter>=1.0.0",
+  "py3dmol>=2.0.4",
+  "tqdm>=4.66.2",
+]
+
+[project.optional-dependencies]
+extras = [
+  "plotly>=5.22.0",
+  "seaborn>=0.13.1",
+  "bokeh"
+]
+
+dev = [
+  "pytest>=5.2",
+  "pytest-cov>=4.1.0",
+  "hypothesis", # ==6.127.3
+  "tox>=4.10.0",
+  "pre-commit>=2.20.0",
+  "ruff==0.4.1",
+  "mypy>=1.9.0",
+  "pandas-stubs",
+  "scipy-stubs",
+  "types-tqdm",
+]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["shnitsel*"]
+
+[tool.ruff]
+include = ["shnitsel/**/*.py"]
+
+[tool.ruff.format]
+quote-style = "preserve"
+
+[tool.ruff.lint]
+ignore = [
+  "E731", # Do not assign a `lambda` expression, use a `def`
+  "E741", # Ambiguous variable name
+  "E70",  # Multiple statements on one line
+]
+
+[tool.ruff.lint.flake8-annotations]
+ignore-fully-untyped = true
+
+[tool.mypy]
+packages = "shnitsel"
+warn_return_any = true
+warn_unused_configs = true
+check_untyped_defs = true
+pretty = true
+color_output = true
+error_summary = true
+
+[[tool.mypy.overrides]]
+module = [
+  "scipy",
+  "sklearn.*"
+]
+ignore_missing_imports = true
+
+[tool.tox]
+min_version = "4.19"
+env_list = ["py313", "py312", "py311", "py310"]
+
+[tool.tox.testenv]
+deps = ["pytest"]
+commands = [
+  ["python", "--version"],
+  ["pytest"]
+]

shnitsel_tools-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

shnitsel_tools-0.0.1/shnitsel/__init__.py

@@ -0,0 +1,14 @@
+from shnitsel import (
+    core as core,
+    plot as plot,
+)
+from shnitsel.core import (
+    parse as parse,
+    postprocess as postprocess,
+    xrhelpers as xrhelpers,
+)
+from shnitsel.core.xrhelpers import open_frames as open_frames
+from shnitsel.core.parse import read_trajs as read_trajs
+from shnitsel.core.ase import read_ase as read_ase
+
+__all__ = ['plot', 'parse', 'open_frames', 'read_trajs', 'read_ase']

shnitsel_tools-0.0.1/shnitsel/core/__init__.py

@@ -0,0 +1,39 @@
+from . import (
+    ase,
+    filter_unphysical,
+    parse,
+    plot,
+    postprocess,
+    xrhelpers,
+)
+
+from .xrhelpers import (
+    open_frames as open_frames,
+    save_frames as save_frames,
+)
+from .postprocess import (
+    dihedral as dihedral,
+    get_per_state as get_per_state,
+    get_inter_state as get_inter_state,
+    assign_fosc as assign_fosc,
+)
+
+from .plot import pca_biplot
+from .plot.spectra3d import spectra_all_times as spectra_all_times
+
+__all__ = [
+    'ase',
+    'parse',
+    'postprocess',
+    'xrhelpers',
+    'filter_unphysical',
+    'pca_biplot',
+    'plot',
+    'open_frames',
+    'save_frames',
+    'dihedral',
+    'get_per_state',
+    'get_inter_state',
+    'assign_fosc',
+    'spectra_all_times',
+]

shnitsel_tools-0.0.1/shnitsel/core/ase.py

@@ -0,0 +1,148 @@
+import os
+from typing import Collection
+
+from ase import Atoms
+from ase.db import connect
+import numpy as np
+import xarray as xr
+
+
+def _prepare_for_write(frames: xr.Dataset) -> xr.Dataset:
+    # Recombine permanent and transition dipoles, as schnetpack expects
+    dipoles: np.ndarray | xr.DataArray | None = None
+    frames = frames.copy(deep=False)
+    if 'dipoles' in frames:
+        dipoles = frames['dipoles']
+    elif 'dip_perm' in frames and 'dip_trans' in frames:
+        dip_perm = frames['dip_perm'].transpose('frame', 'state', 'direction').data
+        dip_trans = (
+            frames['dip_trans'].transpose('frame', 'statecomb', 'direction').data
+        )
+        dipoles = np.concat((dip_perm, dip_trans.data), axis=1)
+        del frames['dip_perm'], frames['dip_trans']
+    elif 'dip_perm' in frames:
+        dipoles = frames['dip_perm']
+        del frames['dip_perm']
+    elif 'dip_trans' in frames:
+        dipoles = frames['dip_trans']
+        del frames['dip_trans']
+
+    if dipoles is not None:
+        frames['dipoles'] = ['frame', 'state_or_statecomb', 'direction'], dipoles
+
+    return frames
+
+
+def write_ase(
+    frames: xr.Dataset,
+    db_path: str,
+    kind: str | None,
+    keys: Collection | None = None,
+    preprocess: bool = True,
+):
+    if preprocess:
+        frames = _prepare_for_write(frames)
+
+    statedims = ['state', 'statecomb', 'state_or_statecomb']
+    if kind == 'schnet':
+        order = ['frame', *statedims, 'atom', 'direction']
+        frames = frames.transpose(*order, missing_dims='ignore')
+    elif kind == 'spainn':
+        frames['energy'] = frames['energy'].expand_dims('tmp', axis=1)
+        order = ['frame', 'tmp', 'atom', *statedims, 'direction']
+        frames = frames.transpose(*order, missing_dims='ignore')
+    elif kind is None:
+        # leave the axis orders as they are
+        pass
+    else:
+        raise ValueError(
+            f"'kind' should be one of 'schnet', 'spainn' or None, not '{kind}'"
+        )
+
+    if os.path.exists(db_path):
+        os.remove(db_path)
+
+    if not keys:
+        keys = frames.data_vars.keys()
+    keys = set(frames.data_vars).intersection(keys).difference({'atNames'})
+
+    with connect(db_path, type='db') as db:
+        for i, frame in frames.groupby('frame'):
+            frame = frame.squeeze('frame')
+            db.write(
+                Atoms(symbols=frame['atNames'].data, positions=frame['atXYZ']),
+                data={k: frame[k].data for k in keys},
+            )
+
+
+def read_ase(db_path: str, kind: str):
+    """Reads an ASE DB containing data in the SPaiNN or SchNet format
+
+    Parameters
+    ----------
+    db_path
+        Path to the database
+    kind
+        Must be one of 'spainn' or 'schnet'; determines interpretation of array shapes
+
+    Returns
+    -------
+        An `xr.Dataset` of frames
+
+    Raises
+    ------
+    ValueError
+        If `kind` is not one of 'spainn' or 'schnet'
+    FileNotFoundError
+        If `db_path` is not a file
+    """
+    if kind == 'schnet':
+        shapes = {
+            'energy': ['frame', 'state'],
+            'forces': ['frame', 'state', 'atom', 'direction'],
+            'nacs': ['frame', 'statecomb', 'atom', 'direction'],
+            'dipoles': ['frame', 'state_or_statecomb', 'direction'],
+        }
+    elif kind == 'spainn':
+        shapes = {
+            'energy': ['frame', 'tmp', 'state'],  # Note the extra dim, removed below
+            'forces': ['frame', 'atom', 'state', 'direction'],
+            'nacs': ['frame', 'atom', 'statecomb', 'direction'],
+            'dipoles': ['frame', 'state_or_statecomb', 'direction'],
+        }
+    else:
+        raise ValueError(f"'kind' should be one of 'schnet' or 'spainn', not '{kind}'")
+
+    if not os.path.isfile(db_path):
+        raise FileNotFoundError(db_path)
+
+    with connect(db_path) as db:
+        data_vars = {}
+        for name, dims in shapes.items():
+            try:
+                data = np.stack([row.data[name] for row in db.select()])
+                data_vars[name] = dims, data
+            except KeyError:
+                pass
+
+        atXYZ = np.stack([row.positions for row in db.select()])
+        data_vars['atXYZ'] = ['frame', 'atom', 'direction'], atXYZ
+        atNames = ['atom'], next(db.select()).symbols
+
+    if 'dipoles' in data_vars:
+        nstates = data_vars['energy'][1].shape[1]
+
+        dipoles = data_vars['dipoles'][1]
+        dip_perm = dipoles[:, :nstates, :]
+        dip_trans = dipoles[:, nstates:, :]
+        del data_vars['dipoles']
+
+        data_vars['dip_perm'] = ['frame', 'state', 'direction'], dip_perm
+        data_vars['dip_trans'] = ['frame', 'statecomb', 'direction'], dip_trans
+
+    frames = xr.Dataset(data_vars).assign_coords(atNames=atNames)
+    if kind == 'spainn':
+        assert 'tmp' in frames.dims
+        frames = frames.squeeze('tmp')
+
+    return frames