"""
A module for writing and saving a H5MD file.
"""
from pathlib import Path
import h5py
import numpy as np
import periodictable
from MDMC.common import units
from MDMC.trajectory_analysis.compact_trajectory import CompactTrajectory
#: Generic H5MD information like, version of H5MD and any H5MD `modules` used.
H5MD_DATA = {
"creator_name": "MDMC",
"creator_version": [0, 2],
"h5md_version": [1, 1],
"module_name": ["units"],
"module_version": [[1, 0]],
"loc": "h5md",
}
ROOT_TRAJECTORY = "particles/all"
[docs]
def create_empty_groups(open_file: h5py.File, groups: list[str]):
"""
Create all groups and subgroups that make up the main base structure of the file format.
Parameters
----------
open_file : h5py.File
A pre-opened file that is being written to
groups : list[str]
A list of groups that are being created
"""
for group_name in groups:
open_file.create_group(group_name)
[docs]
def create_simulation_data(
open_file: h5py.File,
group_name: str,
value: int | float | str | np.ndarray,
unit: str = None,
time_increment: int = None,
step_increment: int = None,
time_offset: int = None,
step_offset: int = None,
):
"""
Store data about the simulation.
Parameters
----------
open_file : h5py.File
A pre-opened file that the data is being written into
group_name : str
The group name for the data being stored
value : int | float | str | np.ndarray
The data being stored in the H5MD file
unit : str, optional
The units for the data being stored. Default is None.
time_increment : int, optional
The increment in time between each simulation step, by default None
step_increment : int, optional
The increment in between steps in a simulation, by default None
time_offset : int, optional
The offset that the time starts at, by default None
step_offset : int, optional
The offset what is where the steps start at, by default None
"""
group = open_file[ROOT_TRAJECTORY]
if time_increment is None:
subdata = group.create_dataset(group_name, data=value)
else:
subgroup = group.create_group(group_name)
subdata = subgroup.create_dataset("value", data=value)
time_link = group.visit(lambda name: name if "time" in name else None)
step_link = group.visit(lambda name: name if "step" in name else None)
if time_link is not None:
time_data = subgroup.create_dataset("time", data=group[time_link])
step_data = subgroup.create_dataset("step", data=group[step_link])
else:
time_data = subgroup.create_dataset("time", data=time_increment)
step_data = subgroup.create_dataset("step", data=step_increment)
time_data.attrs["offset"] = time_offset
time_data.attrs["unit"] = str(units.SYSTEM["TIME"])
step_data.attrs["offset"] = step_offset
if unit is not None:
subdata.attrs["unit"] = str(unit)
[docs]
def create_box_data(open_file: h5py.File, trajectory: CompactTrajectory):
"""
Create the box group and add all attributes associated with this group.
Parameters
----------
trajectory : CompactTrajectory
The compact trajectory the file is being built from
open_file : h5py.File
A pre-opened file that the data is being written into
"""
box_group = open_file[f"{ROOT_TRAJECTORY}/box"]
box_group.attrs["dimensions"] = len(trajectory.dimensions)
if trajectory.is_fixedbox:
boundary = ["periodic" for _ in trajectory.dimensions] # MDMC assumes periodic
box_group.attrs["boundary"] = boundary
[docs]
def create_parameter_data(open_file: h5py.File, data: np.array):
"""
Create and store the remaining parameter data in the H5MD file.
Parameters
----------
open_file : h5py.File
A pre-opened file that the data is being written into
data : np.array
The data being stored
"""
parameters = open_file["parameters"]
parameters.create_dataset("atom_symbols", data=data)
[docs]
def write_H5MD(
trajectory: CompactTrajectory,
filename: str = "trajectory.h5",
*,
timestamp: str,
file_loc: Path = Path("."),
**settings,
):
"""
Write a CompactTrajectory to a H5MD file.
In MDMC, an H5MD trajectory File is built from a
:class:`~MDMC.trajectory_analysis.compact_trajectory.CompactTrajectory`.
Once a ``CompactTrajectory`` is generated, an H5MD trajectory file can be
created by passing it to the :func:`~MDMC.writers.H5MD_build.build_full` function.
Once executed the a H5MD (``.h5``) file is created with the name
"trajectory<timestamp>"
Parameters
----------
trajectory : CompactTrajectory
The compact trajectory the file is being built from.
filename : str, optional
The name of the H5MD file,
standards suggest it ends with `.h5`, by default "trajectory.h5".
timestamp : bool, optional
If true adds time timestamp to file name, by default True.
file_loc : Path, optional
The directory where the H5MD file should be stored, by default Path('.').
"""
if trajectory.time is not None:
time_increment = trajectory.time[1] - trajectory.time[0]
time_offset = trajectory.time[0]
step_increment = 1
step_offset = 0
else:
time_increment = None
time_offset = None
step_increment = None
step_offset = None
file_path_name = Path(file_loc, f"{filename}{timestamp}_traj").with_suffix('.h5')
if not settings.get("creator_name") or not settings.get("creator_email"):
raise ValueError("No creator_name or creator_email provided.")
with h5py.File(file_path_name, "w") as file:
no_data_groups = [
"particles",
ROOT_TRAJECTORY,
"h5md",
f"{ROOT_TRAJECTORY}/box",
"parameters",
]
create_empty_groups(file, no_data_groups)
create_metadata_group(
file, creator_name=settings["creator_name"], creator_email=settings["creator_email"],
)
charge = trajectory.atom_charges
for count, value in enumerate(charge):
if value is None:
charge[count] = 0.0
charge = charge.astype(float)
species = []
symbol_list = []
for element in trajectory.element_list:
if "-" not in element:
species.append(getattr(periodictable, element).number)
symbol_list.append(element)
else:
isotope, symbol = element.split("-")
species.append(getattr(periodictable, symbol)[int(isotope)].number)
symbol_list.append(f"{symbol}{isotope}")
simulation_data = {
"species": [species],
"charge": [charge, units.SYSTEM["CHARGE"]],
"mass": [trajectory.atom_masses, units.SYSTEM["MASS"]],
"position": [
trajectory.position,
trajectory.position_unit,
time_increment,
step_increment,
time_offset,
step_offset,
],
"box/edges": [
trajectory.dimensions,
trajectory.position_unit,
time_increment,
step_increment,
time_offset,
step_offset,
],
}
if trajectory.has_velocity:
simulation_data["velocity"] = [
trajectory.velocity,
trajectory.velocity_unit,
time_increment,
step_increment,
time_offset,
step_offset,
]
for grp_name, data in simulation_data.items():
create_simulation_data(file, grp_name, *data)
create_box_data(file, trajectory)
atom_symbols_data = np.array(symbol_list, dtype=object)
create_parameter_data(file, atom_symbols_data)