# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
"""Utilities for loading example data for POWGEN."""
from pathlib import Path
import scipp as sc
import scippnexus as snx
from ess.powder.types import (
AccumulatedProtonCharge,
CalibrationData,
CalibrationFilename,
DetectorBankSizes,
Filename,
Position,
ProtonCharge,
RawDataAndMetadata,
RunType,
TofDetector,
)
from ess.reduce.data import Entry, make_registry
_registry = make_registry(
"ess/powgen",
version="1",
files={
# Files loadable with Mantid
"PG3_4844_event.nxs": "md5:d5ae38871d0a09a28ae01f85d969de1e",
"PG3_4866_event.nxs": "md5:3d543bc6a646e622b3f4542bc3435e7e",
"PG3_5226_event.nxs": "md5:58b386ebdfeb728d34fd3ba00a2d4f1e",
"PG3_FERNS_d4832_2011_08_24.cal": "md5:c181221ebef9fcf30114954268c7a6b6",
# Zipped Scipp HDF5 files
"PG3_4844_event.zip": Entry(
alg="md5", chk="a644c74f5e740385469b67431b690a3e", unzip=True
),
"PG3_4866_event.zip": Entry(
alg="md5", chk="5bc49def987f0faeb212a406b92b548e", unzip=True
),
"PG3_FERNS_d4832_2011_08_24.zip": Entry(
alg="md5", chk="0fef4ed5f61465eaaa3f87a18f5bb80d", unzip=True
),
"PG3_FERNS_d4832_2011_08_24_spectrum.h5": "md5:7aee0b40deee22d57e21558baa7a6a1a", # noqa: E501
# Smaller files for unit tests
"TEST_PG3_4844_event.h5": "md5:03ea1018c825b0d90a67b4fc7932ea3d",
"TEST_PG3_4866_event.h5": "md5:4a8df454871cec7de9ac624ebdc97095",
"TEST_PG3_FERNS_d4832_2011_08_24_spectrum.h5": "md5:b577a054e9aa7b372df79bf4489947d0", # noqa: E501
},
)
[docs]
def powgen_tutorial_mantid_sample_file() -> Path:
return _registry.get_path("PG3_4844_event.nxs")
[docs]
def powgen_tutorial_mantid_vanadium_file() -> Path:
return _registry.get_path("PG3_4866_event.nxs")
[docs]
def powgen_tutorial_mantid_empty_instrument_file() -> Path:
return _registry.get_path("PG3_5226_event.nxs")
[docs]
def powgen_tutorial_mantid_calibration_file() -> Path:
return _registry.get_path("PG3_FERNS_d4832_2011_08_24.cal")
[docs]
def powgen_tutorial_sample_file(*, small: bool = False) -> Path:
"""
Return the path to the POWGEN sample file.
Parameters
----------
small:
If True, return a smaller file for unit tests.
The small version of the file was created using the following code, which keeps
only 7 columns out of 154 (154 / 7 = 22):
.. code-block:: python
import scipp as sc
fname = 'PG3_4844_event.h5'
dg = sc.io.load_hdf5(fname)
sizes = {"bank": 23, "column": 154, "row": 7}
def foldme(x, dim):
return x.fold(dim=dim, sizes=sizes)['column', ::22].flatten(
dims=list(sizes.keys()), to=dim)
small = sc.DataGroup({
'data': foldme(dg['data'], 'spectrum'),
'detector_info': sc.Dataset(
coords={key: foldme(c, 'detector')
for key, c in dg['detector_info'].coords.items()})
})
sc.io.save_hdf5(small, 'TEST_PG3_4844_event.h5')
"""
prefix = "TEST_" if small else ""
ext = ".h5" if small else ".zip"
return _registry.get_path(f"{prefix}PG3_4844_event{ext}")
[docs]
def powgen_tutorial_vanadium_file(*, small: bool = False) -> Path:
"""
Return the path to the POWGEN vanadium file.
Parameters
----------
small:
If True, return a smaller file for unit tests.
The small version of the file was created using the following code, which keeps
only 7 columns out of 154 (154 / 7 = 22):
.. code-block:: python
import scipp as sc
fname = 'PG3_4866_event.h5'
dg = sc.io.load_hdf5(fname)
sizes = {"bank": 23, "column": 154, "row": 7}
def foldme(x, dim):
return x.fold(dim=dim, sizes=sizes)['column', ::22].flatten(
dims=list(sizes.keys()), to=dim)
small = sc.DataGroup({
'data': foldme(dg['data'], 'spectrum'),
'proton_charge': dg['proton_charge']['pulse_time', ::10]
})
sc.io.save_hdf5(small, 'TEST_PG3_4866_event.h5')
"""
prefix = "TEST_" if small else ""
ext = ".h5" if small else ".zip"
return _registry.get_path(f"{prefix}PG3_4866_event{ext}")
[docs]
def powgen_tutorial_calibration_file(*, small: bool = False) -> Path:
"""
Return the path to the POWGEN calibration file.
Parameters
----------
small:
If True, return a smaller file for unit tests.
The small version of the file was created using the following code, which keeps
only 7 columns out of 154 (154 / 7 = 22):
.. code-block:: python
import scipp as sc
fname = 'PG3_FERNS_d4832_2011_08_24_spectrum.h5'
dg = sc.io.load_hdf5(fname)
sizes = {"bank": 23, "column": 154, "row": 7}
def foldme(x, dim):
return x.fold(dim=dim, sizes=sizes)['column', ::22].flatten(
dims=list(sizes.keys()), to=dim)
small = sc.Dataset(
data={k: foldme(a, 'spectrum') for k, a in ds.items()},
coords={k: foldme(c, 'spectrum') for k, c in ds.coords.items()}
)
sc.io.save_hdf5(small, 'TEST_PG3_FERNS_d4832_2011_08_24_spectrum.h5')
"""
prefix = "TEST_" if small else ""
return _registry.get_path(f"{prefix}PG3_FERNS_d4832_2011_08_24_spectrum.h5")
[docs]
def pooch_load(filename: Filename[RunType]) -> RawDataAndMetadata[RunType]:
"""Load a file with pooch.
If the file is a zip archive, it is extracted and a path to the contained
file is returned.
The loaded data holds both the events and any metadata from the file.
"""
return RawDataAndMetadata[RunType](sc.io.load_hdf5(filename))
[docs]
def pooch_load_calibration(
filename: CalibrationFilename,
detector_dimensions: DetectorBankSizes,
) -> CalibrationData:
"""Load the calibration data for the POWGEN test data."""
if filename is None:
return CalibrationFilename(None)
ds = sc.io.load_hdf5(filename)
ds = sc.Dataset(
{
key: da.fold(dim='spectrum', sizes=detector_dimensions)
for key, da in ds.items()
}
)
return CalibrationData(ds)
[docs]
def source_position(dg: RawDataAndMetadata[RunType]) -> Position[snx.NXsource, RunType]:
"""Return the source position from a loaded data group."""
return Position[snx.NXsource, RunType](dg["data"].coords["source_position"])
[docs]
def sample_position(dg: RawDataAndMetadata[RunType]) -> Position[snx.NXsample, RunType]:
"""Return the sample position from a loaded data group."""
return Position[snx.NXsample, RunType](dg["data"].coords["sample_position"])
providers = (
pooch_load,
pooch_load_calibration,
extract_accumulated_proton_charge,
extract_proton_charge,
extract_raw_data,
sample_position,
source_position,
)
"""Sciline Providers for loading POWGEN data."""
