from pathlib import Path
import h5py
import numpy as np
from discopat.core import Frame, Movie
from discopat.repositories.local import DATA_DIR_PATH, LocalRepository
[docs]
class HDF5Repository(LocalRepository):
def __init__(self, name: str):
super().__init__(name)
self._input_directory_path = DATA_DIR_PATH / self.name
[docs]
def read(self, content_path: str or Path) -> Movie:
file_path = (
self._input_directory_path / content_path / "data_TOKAM_run_00.h5"
)
with h5py.File(file_path) as f:
density = np.array(f["density"])
num_frames = density.shape[0]
width = density.shape[2]
height = density.shape[1]
return Movie(
name=f"{content_path}_density",
frames=[
Frame(
name=str(i),
width=width,
height=height,
annotations=[],
image_array=density[i],
)
for i in range(num_frames)
],
tracks=[],
)
[docs]
def write(self, content_path, content) -> None:
pass
[docs]
class HDF5PotentialRepository(LocalRepository):
def __init__(self, name: str):
super().__init__(name)
self._input_directory_path = DATA_DIR_PATH / self.name
[docs]
def read(self, content_path: str or Path) -> Movie:
file_path = (
self._input_directory_path / content_path / "data_TOKAM_run_00.h5"
)
with h5py.File(file_path) as f:
potential = np.array(f["potential"])
num_frames = potential.shape[0]
width = potential.shape[2]
height = potential.shape[1]
return Movie(
name=f"{content_path}_potential",
frames=[
Frame(
name=str(i),
width=width,
height=height,
annotations=[],
image_array=potential[i],
)
for i in range(num_frames)
],
tracks=[],
)
[docs]
def write(self, content_path, content) -> None:
pass