#!/usr/bin/env python3
import h5py
import numpy as np
import os
import json
import argparse
# ENDF interpolation codes
INTERPOLATION_TYPES = {
0: "none",
1: "histogram",
2: "linear-linear",
3: "linear-log",
4: "log-linear",
5: "log-log"
}
class CompactJSONEncoder(json.JSONEncoder):
"""A JSON Encoder that puts arrays of primitives on a single line."""
def __init__(self, *args, **kwargs):
# Default to 2-space indentation if not specified
if kwargs.get("indent") is None:
kwargs["indent"] = 2
super().__init__(*args, **kwargs)
self.indentation_level = 0
def default(self, obj):
"""Handle NumPy types"""
if isinstance(obj, np.ndarray):
return obj.tolist()
if isinstance(obj, np.integer):
return int(obj)
if isinstance(obj, np.floating):
return float(obj)
return json.JSONEncoder.default(self, obj)
def encode(self, data):
"""Encode JSON objects with special handling for arrays of primitives."""
if isinstance(data, (list, tuple)):
# Put arrays of primitives on a single line
if all(isinstance(x, (int, float, bool, str, type(None))) for x in data):
return "[" + ", ".join(json.dumps(el) for el in data) + "]"
# For other containers, use standard formatting
return self._encode_container(data, "[", "]")
if isinstance(data, dict):
if not data:
return "{}"
# Ensure keys are strings
data = {str(k) if k is not None else "null": v for k, v in data.items()}
if self.sort_keys:
data = dict(sorted(data.items(), key=lambda x: x[0]))
# Format dictionary
return self._encode_container(data, "{", "}")
# For primitive types, use standard JSON encoding
return json.dumps(
data,
skipkeys=self.skipkeys,
ensure_ascii=self.ensure_ascii,
check_circular=self.check_circular,
allow_nan=self.allow_nan,
sort_keys=self.sort_keys,
separators=(self.item_separator, self.key_separator),
default=self.default,
)
def _encode_container(self, container, open_char, close_char):
"""Encode a container with proper indentation."""
is_dict = isinstance(container, dict)
# Handle very small containers on single line (dicts with 1 item or less)
if is_dict and len(container) <= 1:
if not container:
return "{}"
key, value = next(iter(container.items()))
return f"{{ {self.encode(key)}: {self.encode(value)} }}"
# Format container across multiple lines
self.indentation_level += 1
if is_dict:
items = [f"{self.indent_str}{self.encode(k)}: {self.encode(v)}" for k, v in container.items()]
else:
items = [f"{self.indent_str}{self.encode(el)}" for el in container]
self.indentation_level -= 1
return f"{open_char}\n" + ",\n".join(items) + f"\n{self.indent_str}{close_char}"
def iterencode(self, data, **kwargs):
"""Required to also work with `json.dump`."""
return self.encode(data)
@property
def indent_str(self) -> str:
"""Get the indentation string for the current level."""
if isinstance(self.indent, int):
return " " * (self.indentation_level * self.indent)
elif isinstance(self.indent, str):
return self.indentation_level * self.indent
else:
return ""
def extract_cross_section_data(filename):
"""Extract cross-section data from HDF5 file into a dictionary"""
result = {}
with h5py.File(filename, 'r') as f:
# Get the nuclide name
nuclide_name = list(f.keys())[0]
nuclide_group = f[nuclide_name]
# Basic nuclide info
result['nuclide'] = nuclide_name
if 'Z' in nuclide_group.attrs:
result['protons'] = int(nuclide_group.attrs['Z'])
if 'N' in nuclide_group.attrs:
result['neutrons'] = int(nuclide_group.attrs['N'])
if 'A' in nuclide_group.attrs:
result['nucleons'] = int(nuclide_group.attrs['A'])
# Extract temperature data
result['temperatures'] = {}
# Get list of reactions
reactions_group = nuclide_group['reactions']
# For each temperature (typically only 294K)
for temp_str in ['294K']: # You might want to make this more dynamic
temp_data = {'reactions': {}}
# Extract energy grid (same for all reactions at this temperature)
if 'energy' in nuclide_group:
# Access the energy group first
energy_group = nuclide_group['energy']
# Then access the dataset for this temperature
if temp_str in energy_group:
# Get the actual data from the dataset
energy_data = energy_group[temp_str][()]
temp_data['energy'] = energy_data.tolist()
# Process each reaction
for reaction_name in reactions_group.keys():
# Extract MT number
mt = int(reaction_name.replace('reaction_', ''))
reaction_group = reactions_group[reaction_name]
# Skip if no cross-section data at this temperature
if temp_str not in reaction_group or 'xs' not in reaction_group[temp_str]:
continue
# Get cross-section dataset
xs_dataset = reaction_group[f'{temp_str}/xs']
reaction_data = {
'cross_section': xs_dataset[()].tolist()
}
# Get threshold index if present
if 'threshold_idx' in xs_dataset.attrs:
reaction_data['energy_threshold_index'] = int(xs_dataset.attrs['threshold_idx'])
# Get interpolation type if present
# Check in the product_0/distribution_0/applicability path
found_interpolation = False
if 'product_0' in reaction_group and not found_interpolation:
try:
product_group = reaction_group['product_0']
if 'distribution_0' in product_group:
dist_group = product_group['distribution_0']
if 'applicability' in dist_group and 'interpolation' in dist_group['applicability'].attrs:
interp_attr = dist_group['applicability'].attrs['interpolation']
if isinstance(interp_attr, np.ndarray) and len(interp_attr) > 0:
interp_type = int(interp_attr[0])
if interp_type in INTERPOLATION_TYPES:
reaction_data['interpolation_type'] = INTERPOLATION_TYPES[interp_type]
found_interpolation = True
except Exception as e:
# Just ignore errors when trying to find interpolation
pass
# Add reaction data to temperature
temp_data['reactions'][mt] = reaction_data
# Add temperature data to result
result['temperatures'][temp_str] = temp_data
return result
data = extract_cross_section_data(filename)
# Write JSON file
with open(output_filename, 'w') as f:
json.dump(data, f, cls=CompactJSONEncoder, indent=2)
