diff --git a/examples/c5g7/continious_energy_transient/input.py b/examples/c5g7/continious_energy_transient/input.py new file mode 100644 index 000000000..4c8ca9ad6 --- /dev/null +++ b/examples/c5g7/continious_energy_transient/input.py @@ -0,0 +1,437 @@ +import numpy as np +import math +import mcdc + +# ============================================================================= +# Materials +# ============================================================================= +# Assumption: dev branch accepts nuclide-name strings for CE materials. +# If your dev build uses a different constructor name, swap mcdc.Material -> . + +mat_gt = mcdc.Material( + nuclide_composition={ + "H1": 0.050347844752850625, + "H2": 7.842394716362082e-06, + "O16": 0.025117935412784034, + # "O17": 9.542402714463945e-06, + "O18": 5.03657582849965e-05, + } +) + +mat_uo2 = mcdc.Material( + nuclide_composition={ + "O16": 0.04585265389377734, + # "O17": 1.7419604031574338e-05, + "O18": 9.19424166352541e-05, + "U235": 0.0007217486041189947, + "U238": 0.02224950230720295, + } +) + + +mat_mox43 = mcdc.Material( + nuclide_composition={ + "O16": 0.04589711643122753, + # "O17": 1.743649552488715e-05, + "O18": 9.203157163056531e-05, + "U235": 0.0003750264168772414, + "U238": 0.02262319599228636, + } +) + +mat_mox7 = mcdc.Material( + nuclide_composition={ + "O16": 0.04583036614158277, + # "O17": 1.741113682662514e-05, + "O18": 9.189772587857765e-05, + "U235": 0.0005581382302893396, + "U238": 0.022404154012604437, + } +) + +mat_mox87 = mcdc.Material( + nuclide_composition={ + "O16": 0.04585265389377734, + # "O17": 1.7419604031574338e-05, + "O18": 9.19424166352541e-05, + "U235": 0.0007217486041189947, + "U238": 0.02224950230720295, + } +) + +mat_gt = mcdc.Material( + nuclide_composition={ + "H1": 0.050347844752850625, + "H2": 7.842394716362082e-06, + "O16": 0.025117935412784034, + # "O17": 9.542402714463945e-06, + "O18": 5.03657582849965e-05, + } +) + +mat_fc = mcdc.Material( + nuclide_composition={ + "H1": 0.050347844752850625, + "H2": 7.842394716362082e-06, + "O16": 0.025117935412784034, + # "O17": 9.542402714463945e-06, + "O18": 5.03657582849965e-05, + } +) + +mat_cr = mcdc.Material( + nuclide_composition={ + "Ag107": 0.023523285675833942, + "Ag109": 0.02185429814297804, + "In113": 0.0003421922042655644, + "In115": 0.007651085167039375, + "Cd106": 3.38816276451386e-05, + "Cd108": 2.4166172970990425e-05, + "Cd110": 0.0003393605596264083, + "Cd111": 0.0003482051612205208, + "Cd112": 0.0006561061533306398, + "Cd113": 0.00033274751904988726, + "Cd114": 0.0007825159207295705, + "Cd116": 0.00020443276053837845, + } +) + +mat_mod = mcdc.Material( + nuclide_composition={ + "H1": 0.050347844752850625, + "H2": 7.842394716362082e-06, + "O16": 0.025117935412784034, + # "O17": 9.542402714463945e-06, + "O18": 5.03657582849965e-05, + } +) + +# ============================================================================= +# Pin cells +# ============================================================================= + +pitch = 1.26 +radius = 0.54 +core_height = 128.52 +reflector_thickness = 21.42 + +# Control rod banks fractions +# All out: 0.0 +# All in : 1.0 +cr1 = np.array([1.0, 1.0, 0.89, 1.0]) +cr1_t = np.array([0.0, 10.0, 15.0, 15.0 + 1.0 - cr1[-2]]) + +cr2 = np.array([1.0, 1.0, 0.0, 0.0, 0.8]) +cr2_t = np.array([0.0, 5.0, 10.0, 15.0, 15.8]) + +cr3 = np.array([0.75, 0.75, 1.0]) +cr3_t = np.array([0.0, 15.0, 15.25]) + +cr4 = np.array([1.0, 1.0, 0.5, 0.5, 1.0]) +cr4_t = np.array( + [0.0, 5.0, 5.0 + (cr4[1] - cr4[2]) / 2 * 10, 15.0, 15.0 + 1.0 - cr4[-2]] +) + +# Tips of the control rod banks +cr1_bottom = core_height * (0.5 - cr1) +cr2_bottom = core_height * (0.5 - cr2) +cr3_bottom = core_height * (0.5 - cr3) +cr4_bottom = core_height * (0.5 - cr4) +cr1_top = cr1_bottom + core_height +cr2_top = cr2_bottom + core_height +cr3_top = cr3_bottom + core_height +cr4_top = cr4_bottom + core_height + +# Durations of the moving tips +cr1_durations = cr1_t[1:] - cr1_t[:-1] +cr2_durations = cr2_t[1:] - cr2_t[:-1] +cr3_durations = cr3_t[1:] - cr3_t[:-1] +cr4_durations = cr4_t[1:] - cr4_t[:-1] + +# Velocities of the moving tips +cr1_velocities = np.zeros((len(cr1) - 1, 3)) +cr2_velocities = np.zeros((len(cr2) - 1, 3)) +cr3_velocities = np.zeros((len(cr3) - 1, 3)) +cr4_velocities = np.zeros((len(cr4) - 1, 3)) +cr1_velocities[:, 2] = (cr1_top[1:] - cr1_top[:-1]) / cr1_durations +cr2_velocities[:, 2] = (cr2_top[1:] - cr2_top[:-1]) / cr2_durations +cr3_velocities[:, 2] = (cr3_top[1:] - cr3_top[:-1]) / cr3_durations +cr4_velocities[:, 2] = (cr4_top[1:] - cr4_top[:-1]) / cr4_durations + +# Surfaces +cy = mcdc.Surface.CylinderZ(center=[0.0, 0.0], radius=radius) +# Control rod top and bottom tips +z1_top = mcdc.Surface.PlaneZ(z=cr1_top[0]) +z1_bottom = mcdc.Surface.PlaneZ(z=cr1_bottom[0]) +z2_top = mcdc.Surface.PlaneZ(z=cr2_top[0]) +z2_bottom = mcdc.Surface.PlaneZ(z=cr2_bottom[0]) +z3_top = mcdc.Surface.PlaneZ(z=cr3_top[0]) +z3_bottom = mcdc.Surface.PlaneZ(z=cr3_bottom[0]) +z4_top = mcdc.Surface.PlaneZ(z=cr4_top[0]) +z4_bottom = mcdc.Surface.PlaneZ(z=cr4_bottom[0]) +# Fuel top +# (Bottom is bounded by the universe cell) +zf = mcdc.Surface.PlaneZ(z=0.5 * core_height) + +# Move the control tips +z1_top.move(cr1_velocities, cr1_durations) +z1_bottom.move(cr1_velocities, cr1_durations) +z2_top.move(cr2_velocities, cr2_durations) +z2_bottom.move(cr2_velocities, cr2_durations) +z3_top.move(cr3_velocities, cr3_durations) +z3_bottom.move(cr3_velocities, cr3_durations) +z4_top.move(cr4_velocities, cr4_durations) +z4_bottom.move(cr4_velocities, cr4_durations) + +# Fission chamber pin +fc = mcdc.Cell(-cy, mat_fc) +mod = mcdc.Cell(+cy, mat_mod) +fission_chamber = mcdc.Universe(cells=[fc, mod]) + +# Fuel rods +uo2 = mcdc.Cell(-cy & -zf, mat_uo2) +mox4 = mcdc.Cell(-cy & -zf, mat_mox43) +mox7 = mcdc.Cell(-cy & -zf, mat_mox7) +mox8 = mcdc.Cell(-cy & -zf, mat_mox87) +moda = mcdc.Cell(-cy & +zf, mat_mod) # Water above pin +fuel_uo2 = mcdc.Universe(cells=[uo2, mod, moda]) +fuel_mox43 = mcdc.Universe(cells=[mox4, mod, moda]) +fuel_mox7 = mcdc.Universe(cells=[mox7, mod, moda]) +fuel_mox87 = mcdc.Universe(cells=[mox8, mod, moda]) + +# Control rods and guide tubes +cr1 = mcdc.Cell(-cy & +z1_bottom & -z1_top, mat_cr) +gt1_lower = mcdc.Cell(-cy & -z1_bottom, mat_gt) +gt1_upper = mcdc.Cell(-cy & +z1_top, mat_gt) +# +cr2 = mcdc.Cell(-cy & +z2_bottom & -z2_top, mat_cr) +gt2_lower = mcdc.Cell(-cy & -z2_bottom, mat_gt) +gt2_upper = mcdc.Cell(-cy & +z2_top, mat_gt) +# +cr3 = mcdc.Cell(-cy & +z3_bottom & -z3_top, mat_cr) +gt3_lower = mcdc.Cell(-cy & -z3_bottom, mat_gt) +gt3_upper = mcdc.Cell(-cy & +z3_top, mat_gt) +# +cr4 = mcdc.Cell(-cy & +z4_bottom & -z4_top, mat_cr) +gt4_lower = mcdc.Cell(-cy & -z4_bottom, mat_gt) +gt4_upper = mcdc.Cell(-cy & +z4_top, mat_gt) +# +control_rod1 = mcdc.Universe(cells=[cr1, gt1_lower, gt1_upper, mod]) +control_rod2 = mcdc.Universe(cells=[cr2, gt2_lower, gt2_upper, mod]) +control_rod3 = mcdc.Universe(cells=[cr3, gt3_lower, gt3_upper, mod]) +control_rod4 = mcdc.Universe(cells=[cr4, gt4_lower, gt4_upper, mod]) + +# ============================================================================= +# Fuel lattices +# ============================================================================= + +# UO2 lattice 1 +u = fuel_uo2 +c = control_rod1 +f = fission_chamber +lattice_1 = mcdc.Lattice( + x=[-pitch * 17 / 2, pitch, 17], + y=[-pitch * 17 / 2, pitch, 17], + universes=[ + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, u, u, u, c, u, u, c, u, u, c, u, u, u, u, u], + [u, u, u, c, u, u, u, u, u, u, u, u, u, c, u, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, c, u, u, c, u, u, c, u, u, c, u, u, c, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, c, u, u, c, u, u, f, u, u, c, u, u, c, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, c, u, u, c, u, u, c, u, u, c, u, u, c, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, u, c, u, u, u, u, u, u, u, u, u, c, u, u, u], + [u, u, u, u, u, c, u, u, c, u, u, c, u, u, u, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + ], +) + +# MOX lattice 2 +l = fuel_mox43 +m = fuel_mox7 +n = fuel_mox87 +c = control_rod2 +f = fission_chamber +lattice_2 = mcdc.Lattice( + x=[-pitch * 17 / 2, pitch, 17], + y=[-pitch * 17 / 2, pitch, 17], + universes=[ + [l, l, l, l, l, l, l, l, l, l, l, l, l, l, l, l, l], + [l, m, m, m, m, m, m, m, m, m, m, m, m, m, m, m, l], + [l, m, m, m, m, c, m, m, c, m, m, c, m, m, m, m, l], + [l, m, m, c, m, n, n, n, n, n, n, n, m, c, m, m, l], + [l, m, m, m, n, n, n, n, n, n, n, n, n, m, m, m, l], + [l, m, c, n, n, c, n, n, c, n, n, c, n, n, c, m, l], + [l, m, m, n, n, n, n, n, n, n, n, n, n, n, m, m, l], + [l, m, m, n, n, n, n, n, n, n, n, n, n, n, m, m, l], + [l, m, c, n, n, c, n, n, f, n, n, c, n, n, c, m, l], + [l, m, m, n, n, n, n, n, n, n, n, n, n, n, m, m, l], + [l, m, m, n, n, n, n, n, n, n, n, n, n, n, m, m, l], + [l, m, c, n, n, c, n, n, c, n, n, c, n, n, c, m, l], + [l, m, m, m, n, n, n, n, n, n, n, n, n, m, m, m, l], + [l, m, m, c, m, n, n, n, n, n, n, n, m, c, m, m, l], + [l, m, m, m, m, c, m, m, c, m, m, c, m, m, m, m, l], + [l, m, m, m, m, m, m, m, m, m, m, m, m, m, m, m, l], + [l, l, l, l, l, l, l, l, l, l, l, l, l, l, l, l, l], + ], +) + +# MOX lattice 3 +l = fuel_mox43 +m = fuel_mox7 +n = fuel_mox87 +c = control_rod3 +f = fission_chamber +lattice_3 = mcdc.Lattice( + x=[-pitch * 17 / 2, pitch, 17], + y=[-pitch * 17 / 2, pitch, 17], + universes=[ + [l, l, l, l, l, l, l, l, l, l, l, l, l, l, l, l, l], + [l, m, m, m, m, m, m, m, m, m, m, m, m, m, m, m, l], + [l, m, m, m, m, c, m, m, c, m, m, c, m, m, m, m, l], + [l, m, m, c, m, n, n, n, n, n, n, n, m, c, m, m, l], + [l, m, m, m, n, n, n, n, n, n, n, n, n, m, m, m, l], + [l, m, c, n, n, c, n, n, c, n, n, c, n, n, c, m, l], + [l, m, m, n, n, n, n, n, n, n, n, n, n, n, m, m, l], + [l, m, m, n, n, n, n, n, n, n, n, n, n, n, m, m, l], + [l, m, c, n, n, c, n, n, f, n, n, c, n, n, c, m, l], + [l, m, m, n, n, n, n, n, n, n, n, n, n, n, m, m, l], + [l, m, m, n, n, n, n, n, n, n, n, n, n, n, m, m, l], + [l, m, c, n, n, c, n, n, c, n, n, c, n, n, c, m, l], + [l, m, m, m, n, n, n, n, n, n, n, n, n, m, m, m, l], + [l, m, m, c, m, n, n, n, n, n, n, n, m, c, m, m, l], + [l, m, m, m, m, c, m, m, c, m, m, c, m, m, m, m, l], + [l, m, m, m, m, m, m, m, m, m, m, m, m, m, m, m, l], + [l, l, l, l, l, l, l, l, l, l, l, l, l, l, l, l, l], + ], +) + +# UO2 lattice 4 +u = fuel_uo2 +c = control_rod4 +f = fission_chamber +lattice_4 = mcdc.Lattice( + x=[-pitch * 17 / 2, pitch, 17], + y=[-pitch * 17 / 2, pitch, 17], + universes=[ + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, u, u, u, c, u, u, c, u, u, c, u, u, u, u, u], + [u, u, u, c, u, u, u, u, u, u, u, u, u, c, u, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, c, u, u, c, u, u, c, u, u, c, u, u, c, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, c, u, u, c, u, u, f, u, u, c, u, u, c, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, c, u, u, c, u, u, c, u, u, c, u, u, c, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, u, c, u, u, u, u, u, u, u, u, u, c, u, u, u], + [u, u, u, u, u, c, u, u, c, u, u, c, u, u, u, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + [u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u, u], + ], +) + +# ============================================================================= +# Assemblies and core +# ============================================================================= + +# Surfaces +x0 = mcdc.Surface.PlaneX(x=0.0, boundary_condition="reflective") +x1 = mcdc.Surface.PlaneX(x=pitch * 17) +x2 = mcdc.Surface.PlaneX(x=pitch * 17 * 2) +x3 = mcdc.Surface.PlaneX(x=pitch * 17 * 3, boundary_condition="vacuum") + +y0 = mcdc.Surface.PlaneY(y=-pitch * 17 * 3, boundary_condition="vacuum") +y1 = mcdc.Surface.PlaneY(y=-pitch * 17 * 2) +y2 = mcdc.Surface.PlaneY(y=-pitch * 17) +y3 = mcdc.Surface.PlaneY(y=0.0, boundary_condition="reflective") + +z0 = mcdc.Surface.PlaneZ( + z=-(core_height / 2 + reflector_thickness), boundary_condition="vacuum" +) +z1 = mcdc.Surface.PlaneZ(z=-(core_height / 2)) +z2 = mcdc.Surface.PlaneZ( + z=(core_height / 2 + reflector_thickness), boundary_condition="vacuum" +) + +# Assembly cells +center = np.array([pitch * 17 / 2, -pitch * 17 / 2, 0.0]) +assembly_1 = mcdc.Cell(+x0 & -x1 & +y2 & -y3 & +z1 & -z2, lattice_1, translation=center) + +center += np.array([pitch * 17, 0.0, 0.0]) +assembly_2 = mcdc.Cell(+x1 & -x2 & +y2 & -y3 & +z1 & -z2, lattice_2, translation=center) + +center += np.array([-pitch * 17, -pitch * 17, 0.0]) +assembly_3 = mcdc.Cell(+x0 & -x1 & +y1 & -y2 & +z1 & -z2, lattice_3, translation=center) + +center += np.array([pitch * 17, 0.0, 0.0]) +assembly_4 = mcdc.Cell(+x1 & -x2 & +y1 & -y2 & +z1 & -z2, lattice_4, translation=center) + +# Bottom reflector cell +reflector_bottom = mcdc.Cell(+x0 & -x3 & +y0 & -y3 & +z0 & -z1, mat_mod) + +# Side reflectors +reflector_south = mcdc.Cell(+x0 & -x3 & +y0 & -y1 & +z1 & -z2, mat_mod) +reflector_east = mcdc.Cell(+x2 & -x3 & +y1 & -y3 & +z1 & -z2, mat_mod) + +# Root universe +mcdc.simulation.set_root_universe( + cells=[ + assembly_1, + assembly_2, + assembly_3, + assembly_4, + reflector_bottom, + reflector_south, + reflector_east, + ], +) + +# ============================================================================= +# Set source +# ============================================================================= +# Throughout the active center pin of Assembly four, at highest energy, +# for the first 15 seconds + +source = mcdc.Source( + x=np.array([pitch * 17 * 3 / 2] * 2) + np.array([-pitch / 2, +pitch / 2]), + y=np.array([-pitch * 17 * 3 / 2] * 2) + np.array([-pitch / 2, +pitch / 2]), + z=[-core_height / 2, core_height / 2], + isotropic=True, + energy_group=0, # Highest energy + time=[0.0, 15.0], +) + +# ============================================================================= +# Set tallies, settings, techniques and run MC/DC +# ============================================================================= + +# Tallies +Nt = 100 +Nx = 17 * 2 +Ny = 17 * 2 +Nz = 17 * 6 +t = np.linspace(0.0, 20.0, Nt + 1) +x = np.linspace(0.0, pitch * 17 * 2, Nx + 1) +y = np.linspace(-pitch * 17 * 2, 0.0, Ny + 1) +z = np.linspace(-core_height / 2, core_height / 2, Nz + 1) +mesh = mcdc.MeshStructured(x=x, y=y, z=z) +mcdc.TallyMesh(mesh=mesh, scores=["fission"], time=t) + +# Settings +mcdc.settings.N_particle = 10000 +mcdc.settings.N_batch = 2 +mcdc.settings.active_bank_buffer = 1000 + +# Run +mcdc.run() diff --git a/examples/slab_sens/input.py b/examples/slab_sens/input.py new file mode 100644 index 000000000..af965c33c --- /dev/null +++ b/examples/slab_sens/input.py @@ -0,0 +1,71 @@ +import os +import numpy as np +import mcdc + +# ----------------------------------------------------------------------------- +# Continuous-energy data +# ----------------------------------------------------------------------------- +# You must point to a directory of nuclide HDF5s, e.g.: +# export MCDC_LIB=/path/to/mcdc_ce_library +# +# Nuclide names must match the file prefix in that directory. +# Example filename: H1-293.6K.h5 -> nuclide name "H1" +# +if "MCDC_LIB" not in os.environ: + raise RuntimeError("Set MCDC_LIB to your continuous-energy nuclear data directory") + +# Material: simple scatterer (use a nuclide you actually have in MCDC_LIB) +mat = mcdc.Material(nuclide_composition={"H1": 0.066}) # density units per your library + +# ----------------------------------------------------------------------------- +# Geometry: slab in x, reflective in y/z +# ----------------------------------------------------------------------------- +L = 10.0 +xm = 5.0 + +sx0 = mcdc.Surface.PlaneX(x=0.0, boundary_condition="vacuum") +sxm = mcdc.Surface.PlaneX(x=xm) +sxL = mcdc.Surface.PlaneX(x=L, boundary_condition="vacuum") + +sy0 = mcdc.Surface.PlaneY(y=-1.0, boundary_condition="reflective") +sy1 = mcdc.Surface.PlaneY(y=+1.0, boundary_condition="reflective") +sz0 = mcdc.Surface.PlaneZ(z=-1.0, boundary_condition="reflective") +sz1 = mcdc.Surface.PlaneZ(z=+1.0, boundary_condition="reflective") + +left = mcdc.Cell( + name="left", region=(+sx0 & -sxm & +sy0 & -sy1 & +sz0 & -sz1), fill=mat +) +right = mcdc.Cell( + name="right", region=(+sxm & -sxL & +sy0 & -sy1 & +sz0 & -sz1), fill=mat +) + +# ----------------------------------------------------------------------------- +# Source: isotropic point near left boundary +# ----------------------------------------------------------------------------- +mcdc.Source( + position=[0.1, 0.0, 0.0], + isotropic=True, + energy=2.0e6, # eV +) + +# ----------------------------------------------------------------------------- +# Tallies: flux in each cell (baseline sanity check) +# ----------------------------------------------------------------------------- +mcdc.TallyCell(name="flux_left", cell=left, scores=["flux"]) +mcdc.TallyCell(name="flux_right", cell=right, scores=["flux"]) + +# ----------------------------------------------------------------------------- +# Sensitivity response regions: per-cell resp_cum +# ----------------------------------------------------------------------------- +mcdc.settings.sensitivity_mode = True +mcdc.settings.sensitivity_n_resp = 2 +mcdc.settings.sensitivity_resp_cell_IDs = np.array([left.ID, right.ID], dtype=np.int64) + +# ----------------------------------------------------------------------------- +# Run settings +# ----------------------------------------------------------------------------- +mcdc.settings.N_particle = 20000 +mcdc.settings.N_batch = 10 +mcdc.simulation.implicit_capture() + +mcdc.run() diff --git a/mcdc/code_factory/adapt.py b/mcdc/code_factory/adapt.py index 891e85633..f2c06eb9e 100644 --- a/mcdc/code_factory/adapt.py +++ b/mcdc/code_factory/adapt.py @@ -326,7 +326,11 @@ def toggle_inner(func): def set_toggle(flag, val): - toggle_rosters[flag][0] = val + global toggle_rosters + if flag not in toggle_rosters: + toggle_rosters[flag] = [bool(val), []] + else: + toggle_rosters[flag][0] = bool(val) def eval_toggle(): diff --git a/mcdc/code_factory/numba_objects_generator.py b/mcdc/code_factory/numba_objects_generator.py index 452375486..849979136 100644 --- a/mcdc/code_factory/numba_objects_generator.py +++ b/mcdc/code_factory/numba_objects_generator.py @@ -170,6 +170,29 @@ def generate_numba_objects(simulation): ) } + # ================================================================================== + # Sensitivity (optional): inject per-particle cumulative response vector `resp_cum` + # only when requested. This keeps particle records lean when sensitivity is off. + # + # We add the field to BOTH `particle_data` (used in banks) and `particle` (used in + # transport). The field is a fixed-size vector of length `sensitivity_n_resp`. + # ================================================================================== + sens_on = bool(getattr(simulation.settings, "sensitivity_mode", False)) + if sens_on: + n_resp = int(getattr(simulation.settings, "sensitivity_n_resp", 0)) + if n_resp <= 0: + raise ValueError( + "sensitivity_mode=True but sensitivity_n_resp is not a positive integer" + ) + + resp_hint = parse_type_hint_str(f"Annotated[NDArray[float64], ({n_resp},)]") + annotations["particle_data"]["resp_cum"] = resp_hint + annotations["particle"]["resp_cum"] = resp_hint + + # NEW: response-region cell IDs (length = n_resp) + cell_hint = parse_type_hint_str(f"Annotated[NDArray[int64], ({n_resp},)]") + annotations["settings"]["sensitivity_resp_cell_IDs"] = cell_hint + # ================================================================================== # Set the structures and accessor targets based on the annotations # ================================================================================== diff --git a/mcdc/main.py b/mcdc/main.py index 4e1943562..2c4e30439 100644 --- a/mcdc/main.py +++ b/mcdc/main.py @@ -5,6 +5,7 @@ from mcdc import mcdc_get from mcdc.print_ import print_error, print_structure +import importlib def run(): @@ -35,6 +36,9 @@ def run(): # Preparation # ================================================================================== + # To compile adaptive sensitivities + import mcdc.transport.simulation + # Timer: preparation time_prep_start = MPI.Wtime() @@ -206,11 +210,18 @@ def preparation(): from mcdc.code_factory.numba_objects_generator import make_literals make_literals(simulation) + + MPI.COMM_WORLD.Barrier() + importlib.invalidate_caches() + + import mcdc.transport.literals as literals + + importlib.reload(literals) + mcdc_arr, data = generate_numba_objects(simulation) mcdc = mcdc_arr[0] # Reload mcdc getters and setters - import importlib import mcdc.mcdc_get as mcdc_get import mcdc.mcdc_set as mcdc_set @@ -250,6 +261,10 @@ def preparation(): type_.particle, type_.particle_data, ) + # Optional sensitivity module: keep kernels/data minimal unless requested + adapt.set_toggle( + "sensitivity", bool(getattr(simulation.settings, "sensitivity_mode", False)) + ) adapt.eval_toggle() adapt.target_for(config.target) diff --git a/mcdc/mcdc_get/settings.py b/mcdc/mcdc_get/settings.py index cb06d70e9..b242d3c6b 100644 --- a/mcdc/mcdc_get/settings.py +++ b/mcdc/mcdc_get/settings.py @@ -30,3 +30,32 @@ def census_time_chunk(start, length, settings, data): start += settings["census_time_offset"] end = start + length return data[start:end] + + +@njit +def sensitivity_resp_cell_IDs(index, settings, data): + offset = settings["sensitivity_resp_cell_IDs_offset"] + return data[offset + index] + + +@njit +def sensitivity_resp_cell_IDs_all(settings, data): + start = settings["sensitivity_resp_cell_IDs_offset"] + size = settings["sensitivity_resp_cell_IDs_length"] + end = start + size + return data[start:end] + + +@njit +def sensitivity_resp_cell_IDs_last(settings, data): + start = settings["sensitivity_resp_cell_IDs_offset"] + size = settings["sensitivity_resp_cell_IDs_length"] + end = start + size + return data[end - 1] + + +@njit +def sensitivity_resp_cell_IDs_chunk(start, length, settings, data): + start += settings["sensitivity_resp_cell_IDs_offset"] + end = start + length + return data[start:end] diff --git a/mcdc/mcdc_set/settings.py b/mcdc/mcdc_set/settings.py index 7c876f2c2..b7089611c 100644 --- a/mcdc/mcdc_set/settings.py +++ b/mcdc/mcdc_set/settings.py @@ -30,3 +30,32 @@ def census_time_chunk(start, length, settings, data, value): start += settings["census_time_offset"] end = start + length data[start:end] = value + + +@njit +def sensitivity_resp_cell_IDs(index, settings, data, value): + offset = settings["sensitivity_resp_cell_IDs_offset"] + data[offset + index] = value + + +@njit +def sensitivity_resp_cell_IDs_all(settings, data, value): + start = settings["sensitivity_resp_cell_IDs_offset"] + size = settings["sensitivity_resp_cell_IDs_length"] + end = start + size + data[start:end] = value + + +@njit +def sensitivity_resp_cell_IDs_last(settings, data, value): + start = settings["sensitivity_resp_cell_IDs_offset"] + size = settings["sensitivity_resp_cell_IDs_length"] + end = start + size + data[end - 1] = value + + +@njit +def sensitivity_resp_cell_IDs_chunk(start, length, settings, data, value): + start += settings["sensitivity_resp_cell_IDs_offset"] + end = start + length + data[start:end] = value diff --git a/mcdc/object_/settings.py b/mcdc/object_/settings.py index 62a82eec1..e9c840f64 100644 --- a/mcdc/object_/settings.py +++ b/mcdc/object_/settings.py @@ -60,6 +60,13 @@ class Settings(ObjectSingleton): source_bank_buffer_ratio: float = 2.0 future_bank_buffer_ratio: float = 1.5 + # Sensitivity analysis (optional) + sensitivity_mode: bool = False + sensitivity_n_resp: int = 0 + sensitivity_resp_cell_IDs: NDArray[np.int64] = field( + default_factory=lambda: np.zeros(0, dtype=np.int64) + ) + # Portability target_gpu: bool = False diff --git a/mcdc/transport/particle.py b/mcdc/transport/particle.py index 5a01b02b2..2186354c0 100644 --- a/mcdc/transport/particle.py +++ b/mcdc/transport/particle.py @@ -5,6 +5,22 @@ import mcdc.transport.physics as physics import mcdc.transport.rng as rng +import mcdc.code_factory.adapt as adapt + + +@adapt.toggle("sensitivity") +def _copy_resp_cum(target_particle_container, source_particle_container): + """Copy per-particle response accumulator (sensitivity mode only).""" + target = target_particle_container[0] + source = source_particle_container[0] + target["resp_cum"][:] = source["resp_cum"][:] + + +@adapt.toggle("sensitivity") +def _reset_resp_cum(particle_container): + """Reset per-particle response accumulator for a new child history.""" + particle_container[0]["resp_cum"][:] = 0.0 + @njit def move(particle_container, distance, mcdc, data): @@ -35,6 +51,8 @@ def copy(target_particle_container, source_particle_container): target_particle["particle_type"] = source_particle["particle_type"] target_particle["rng_seed"] = source_particle["rng_seed"] + _copy_resp_cum(target_particle_container, source_particle_container) + @njit def copy_as_child(child_particle_container, parent_particle_container): @@ -43,6 +61,8 @@ def copy_as_child(child_particle_container, parent_particle_container): copy(child_particle_container, parent_particle_container) + _reset_resp_cum(child_particle_container) + # Set child RNG seed based of the parent parent_seed = parent_particle["rng_seed"] child_particle["rng_seed"] = rng.split_seed(parent_seed, rng.SEED_SPLIT_PARTICLE) diff --git a/mcdc/transport/particle_bank.py b/mcdc/transport/particle_bank.py index a989db4d3..cda6e2a47 100644 --- a/mcdc/transport/particle_bank.py +++ b/mcdc/transport/particle_bank.py @@ -20,6 +20,12 @@ from mcdc.print_ import print_error +@adapt.toggle("sensitivity") +def _unbank_resp_cum(P_arr, P_rec): + """Copy resp_cum from a bank particle record into a live particle (sens only).""" + P_arr[0]["resp_cum"][:] = P_rec["resp_cum"][:] + + # ============================================================================= # Particle bank operations # ============================================================================= @@ -81,6 +87,8 @@ def get_particle(P_arr, bank, mcdc): P["particle_type"] = P_rec["particle_type"] P["rng_seed"] = P_rec["rng_seed"] + _unbank_resp_cum(P_arr, P_rec) + # Set default IDs and event P["alive"] = True P["material_ID"] = -1 diff --git a/mcdc/transport/sensitivity.py b/mcdc/transport/sensitivity.py new file mode 100644 index 000000000..c79300a79 --- /dev/null +++ b/mcdc/transport/sensitivity.py @@ -0,0 +1,24 @@ +import mcdc.code_factory.adapt as adapt + + +@adapt.toggle("sensitivity") +def score_resp_cum_tracklength(particle_container, distance, mcdc, data): + """ + Sensitivity-only track-length scoring into particle["resp_cum"][k], + where k indexes the user-defined response regions. + + Current definition: response regions are specified by cell_ID membership. + """ + # Keep consistent with other tallying behavior + if not mcdc["cycle_active"]: + return + + P = particle_container[0] + cell_id = P["cell_ID"] + + resp_cells = mcdc["settings"]["sensitivity_resp_cell_IDs"] + # Find which response region this cell corresponds to (small N_resp => linear scan OK) + for k in range(resp_cells.shape[0]): + if cell_id == resp_cells[k]: + P["resp_cum"][k] += P["w"] * distance + break diff --git a/mcdc/transport/simulation.py b/mcdc/transport/simulation.py index 0391a72cb..3c40d6c9c 100644 --- a/mcdc/transport/simulation.py +++ b/mcdc/transport/simulation.py @@ -18,6 +18,7 @@ import mcdc.transport.rng as rng import mcdc.transport.tally as tally_module import mcdc.transport.technique as technique +import mcdc.transport.sensitivity as sensitivity from mcdc.constant import * from mcdc.print_ import ( @@ -622,6 +623,10 @@ def move_to_event(particle_container, mcdc, data): if settings["eigenvalue_mode"]: tally_module.score.eigenvalue_tally(particle_container, distance, mcdc, data) + if settings["sensitivity_mode"]: + # Sensitivity response accumulator (track-length) + sensitivity.score_resp_cum_tracklength(particle_container, distance, mcdc, data) + # Move particle particle_module.move(particle_container, distance, mcdc, data) diff --git a/mcdc/transport/source.py b/mcdc/transport/source.py index 244cccff3..7aba41d77 100644 --- a/mcdc/transport/source.py +++ b/mcdc/transport/source.py @@ -15,6 +15,14 @@ ) from mcdc.transport.util import find_bin +import mcdc.code_factory.adapt as adapt + + +@adapt.toggle("sensitivity") +def _init_resp_cum(P_rec_arr): + """Initialize per-particle response accumulator for a new source history.""" + P_rec_arr[0]["resp_cum"][:] = 0.0 + @njit def source_particle(P_rec_arr, seed, mcdc, data): @@ -110,3 +118,5 @@ def source_particle(P_rec_arr, seed, mcdc, data): P_rec["E"] = E P_rec["w"] = 1.0 P_rec["particle_type"] = source["particle_type"] + + _init_resp_cum(P_rec_arr) diff --git a/pyproject.toml b/pyproject.toml index 8c2c73d21..bd329ed80 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,6 +82,7 @@ dev = [ "black", "pre-commit", "pytest", + "tqdm", # for continuos energy processing ] [project.urls] diff --git a/tools/data_library_generator/generate.py b/tools/data_library_generator/generate.py index 3e595a817..dedf3630a 100644 --- a/tools/data_library_generator/generate.py +++ b/tools/data_library_generator/generate.py @@ -1,4 +1,3 @@ -import ACEtk import argparse import h5py import numpy as np @@ -6,11 +5,32 @@ from tqdm import tqdm -#### - import util from util import print_error, print_note +try: + import ACEtk + import tools as njoy_tools + + # Sanity check: if this is not a compiled extension, it's likely the *wrong* "tools" package + tools_path = getattr(njoy_tools, "__file__", "") + if not tools_path or not tools_path.endswith((".so", ".pyd", ".dylib")): + print_error( + "Imported a Python package named 'tools' that doesn't look like the NJOY tools bindings.\n" + f"Resolved tools at: {tools_path}\n" + "Set PYTHONPATH=$PYTHONPATH:/build/_deps/tools-build/python\n" + ) +except: + print_error( + "ACEtk is not installed\n" + " 1. Follow installation instructions at:\n" + " https://github.com/njoy/ACEtk\n" + " 2. Set\n" + " PYTHONPATH=$PYTHONPATH:/build/python\n" + " PYTHONPATH=$PYTHONPATH:/build/_deps/tools-build/python\n" + ) + + parser = argparse.ArgumentParser(description="MC/DC data generator") parser.add_argument("--rewrite", dest="rewrite", action="store_true", default=False) parser.add_argument("--verbose", dest="verbose", action="store_true", default=False) @@ -57,6 +77,7 @@ disable=verbose, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}{postfix}", ) + for ace_name in pbar: # File header with open(f"{ace_dir}/{ace_name}", "r") as f: @@ -210,7 +231,7 @@ xs0_block = ace_table.principal_cross_section_block xs_block = ace_table.cross_section_block - xs_energy = xs0_block.energies + xs_energy = np.array(xs0_block.energies.to_list()) xs_elastic = xs0_block.elastic cross_sections = xs_block.cross_sections offsets = xs_block.energy_index @@ -320,7 +341,9 @@ # The distributions energy_group = group.create_group(f"MT-{MT:03}/energy_spectrum-1") - util.load_energy_distribution(data, energy_group) + util.load_energy_distribution( + data, energy_group, incident_grid=xs_energy + ) else: N_dist = data.number_distributions @@ -380,7 +403,9 @@ f"MT-{MT:03}/energy_spectrum-{i+1}" ) distribution = data.distribution(i + 1) - util.load_energy_distribution(distribution, energy_group) + util.load_energy_distribution( + data, energy_group, incident_grid=xs_energy + ) # Fissionable zone below if not fissionable: @@ -412,7 +437,7 @@ decay_rates = np.zeros(N_DNP) for i in range(N_DNP): - idx = 1 + 1 + idx = i + 1 data = dnp_block.precursor_group_data(idx) if ( @@ -439,7 +464,7 @@ N_DNP = dnp_block.number_delayed_precursors for i in range(N_DNP): - idx = 1 + 1 + idx = i + 1 data = delayed_spectrum_block.energy_distribution_data(idx) if not isinstance(data, ACEtk.continuous.OutgoingEnergyDistributionData): @@ -448,7 +473,7 @@ energy_group = fission_group.create_group( f"delayed_neutron_precursors/energy_spectrum-{i+1}" ) - util.load_energy_distribution(data, energy_group) + util.load_energy_distribution(data, energy_group, incident_grid=xs_energy) # ================================================================================== # Finalize diff --git a/tools/data_library_generator/util.py b/tools/data_library_generator/util.py index eca313517..9eee0808d 100644 --- a/tools/data_library_generator/util.py +++ b/tools/data_library_generator/util.py @@ -1,6 +1,7 @@ import ACEtk import h5py import numpy as np +import math def print_error(message): @@ -78,6 +79,68 @@ def get_ace_name(Z, A, T, S=None): return f"{ID}{extension}" +def _as_np(x): + """Convert ACEtk views / python lists to numpy arrays reliably.""" + if hasattr(x, "to_list"): + return np.array(x.to_list()) + return np.array(x) + + +def _endf_interp(x, x1, x2, y1, y2, law): + # ENDF interpolation laws: + # 1 = histogram + # 2 = lin-lin + # 3 = lin-log(x) + # 4 = log(y)-lin(x) + # 5 = log(y)-log(x) + if law == 1: + return y1 + if law == 2: + return y1 + (y2 - y1) * (x - x1) / (x2 - x1) + if law == 3: + return y1 + (y2 - y1) * (math.log(x / x1) / math.log(x2 / x1)) + if law == 4: + return math.exp( + math.log(y1) + (math.log(y2) - math.log(y1)) * (x - x1) / (x2 - x1) + ) + if law == 5: + return math.exp( + math.log(y1) + + (math.log(y2) - math.log(y1)) * (math.log(x / x1) / math.log(x2 / x1)) + ) + raise ValueError(f"Unsupported interpolation law: {law}") + + +def _tab1_eval(x, grid, vals, boundaries, interpolants): + """Evaluate a TAB1-like table at x using ENDF interpolation regions.""" + # Clamp outside range + if x <= grid[0]: + return vals[0] + if x >= grid[-1]: + return vals[-1] + + i = int(np.searchsorted(grid, x) - 1) + i = max(0, min(i, len(grid) - 2)) + + # boundaries are 1-based point indices for region ends + law = int(interpolants[-1]) + for k, end_pt_1based in enumerate(boundaries): + if i <= (int(end_pt_1based) - 2): + law = int(interpolants[k]) + break + + x1, x2 = float(grid[i]), float(grid[i + 1]) + y1, y2 = float(vals[i]), float(vals[i + 1]) + + # log-safety fallback + if law in (3, 5) and (x <= 0 or x1 <= 0 or x2 <= 0): + law = 2 + if law in (4, 5) and (y1 <= 0 or y2 <= 0): + law = 2 + + return _endf_interp(float(x), x1, x2, y1, y2, law) + + def load_fission_multiplicity(data, h5_group: h5py.Group): # Polynomial if data.type == 1: @@ -149,11 +212,11 @@ def load_cosine_distribution(data, h5_group: h5py.Group): print_error("Angular distribution is not linearly-iterpolable") -def load_energy_distribution(data, h5_group: h5py.Group): +def load_energy_distribution(data, h5_group: h5py.Group, incident_grid=None): if isinstance(data, ACEtk.continuous.LevelScatteringDistribution): h5_group.attrs["type"] = "level-scattering" - C1 = np.array(data.C1) + C1 = _as_np(data.C1) C1 = h5_group.create_dataset("C1", data=C1) C1.attrs["unit"] = "MeV" @@ -162,21 +225,185 @@ def load_energy_distribution(data, h5_group: h5py.Group): elif isinstance(data, ACEtk.continuous.EvaporationSpectrum): h5_group.attrs["type"] = "evaporation" + # Raw tabulation (what ACE gives you) + energy_raw = _as_np(data.energies) + temperature_raw = _as_np(data.temperatures) + restriction_energy = _as_np(data.restriction_energy) + + # TAB1-like interpolation metadata + boundaries = _as_np(data.interpolation_data.boundaries).astype(int) + interpolants = _as_np(data.interpolation_data.interpolants).astype(int) + + # If non lin-lin, resample temperature onto incident_grid (xs_energy) if not data.interpolation_data.is_linear_linear: + if incident_grid is None: + print_error( + "Evaporation temperature is not linear-linear. " + "Pass incident_grid so it can be resampled." + ) + + xg = np.asarray(incident_grid, dtype=float) + temperature = np.array( + [ + _tab1_eval(x, energy_raw, temperature_raw, boundaries, interpolants) + for x in xg + ], + dtype=float, + ) + energy = xg + h5_group.create_dataset("temperature_interpolation_resampled", data=True) + else: + energy = energy_raw + temperature = temperature_raw + h5_group.create_dataset("temperature_interpolation_resampled", data=False) + + # Save raw + interpolation info (traceability) + h5_group.create_dataset("temperature_energy_grid_raw", data=energy_raw).attrs[ + "unit" + ] = "MeV" + h5_group.create_dataset("temperature_raw", data=temperature_raw).attrs[ + "unit" + ] = "MeV" + h5_group.create_dataset("temperature_interp_boundaries", data=boundaries) + h5_group.create_dataset("temperature_interp_interpolants", data=interpolants) + + # Save the “effective” temperature grid/value used by the generator + h5_group.create_dataset("temperature_energy_grid", data=energy).attrs[ + "unit" + ] = "MeV" + h5_group.create_dataset("temperature", data=temperature).attrs["unit"] = "MeV" + h5_group.create_dataset("restriction_energy", data=restriction_energy).attrs[ + "unit" + ] = "MeV" + + elif isinstance(data, ACEtk.continuous.SimpleMaxwellianFissionSpectrum): + h5_group.attrs["type"] = "maxwellian" + + if all(_as_np(data.interpolation_data.interpolants) == 2): + interpolation = "linear" + elif all(_as_np(data.interpolation_data.interpolants) == 5): + interpolation = "log" + else: print_error( - "Evaporation distribution temperature is not linearly interpolable" + "Unsupported temperature interpolation law in Maxwellian distribution" ) - energy = np.array(data.energies) - temperature = np.array(data.temperatures) - restriction_energy = np.array(data.restriction_energy) + energy = _as_np(data.energies) + temperature = _as_np(data.temperatures) + restriction_energy = _as_np(data.restriction_energy) - dataset = h5_group.create_dataset("temperature_energy_grid", data=energy) - dataset.attrs["unit"] = "MeV" - dataset = h5_group.create_dataset("temperature", data=temperature) - dataset.attrs["unit"] = "MeV" - dataset = h5_group.create_dataset("restriction_energy", data=restriction_energy) - dataset.attrs["unit"] = "MeV" + h5_group.create_dataset("temperature_interpolation", data=interpolation) + h5_group.create_dataset("temperature_energy_grid", data=energy).attrs[ + "unit" + ] = "MeV" + h5_group.create_dataset("temperature", data=temperature).attrs["unit"] = "MeV" + h5_group.create_dataset("restriction_energy", data=restriction_energy).attrs[ + "unit" + ] = "MeV" + + elif isinstance(data, ACEtk.continuous.OutgoingEnergyDistributionData): + h5_group.attrs["type"] = "tabulated" + + if not data.interpolation_data.is_linear_linear: + print_error( + "Non-linearly-interpolated energy distribution is not supported" + ) + + energy = _as_np(data.incident_energies) + h5_group.create_dataset("energy", data=energy).attrs["unit"] = "MeV" + + NE = data.number_incident_energies + offset = np.zeros(NE, dtype=int) + energy_out = [] + pdf = [] + for i in range(NE): + distribution = data.distribution(i + 1) + offset[i] = len(energy_out) + energy_out.extend(distribution.outgoing_energies) + pdf.extend(distribution.pdf) + + energy_out = np.array(energy_out) + pdf = np.array(pdf) + + h5_group.create_dataset("offset", data=offset) + h5_group.create_dataset("value", data=energy_out).attrs["unit"] = "MeV" + h5_group.create_dataset("pdf", data=pdf) + + elif isinstance(data, ACEtk.continuous.KalbachMannDistributionData): + h5_group.attrs["type"] = "kalbach-mann" + + if not data.interpolation_data.is_linear_linear: + print_error("Non-linearly-interpolated kalbach-mann is not supported") + + NE = data.number_incident_energies + energy = _as_np(data.incident_energies) + h5_group.create_dataset("energy", data=energy).attrs["unit"] = "MeV" + + offset = np.zeros(NE, dtype=int) + energy_out = [] + pdf = [] + precompound_factor = [] + angular_slope = [] + for i, distribution in enumerate(data.distributions): + offset[i] = len(pdf) + energy_out.extend(distribution.outgoing_energies) + pdf.extend(distribution.pdf) + precompound_factor.extend(distribution.precompound_fraction_values) + angular_slope.extend(distribution.angular_distribution_slope_values) + + h5_group.create_dataset("offset", data=offset) + h5_group.create_dataset("energy_out", data=np.array(energy_out)).attrs[ + "unit" + ] = "MeV" + h5_group.create_dataset("pdf", data=np.array(pdf)) + h5_group.create_dataset("precompound_factor", data=np.array(precompound_factor)) + h5_group.create_dataset("angular_slope", data=np.array(angular_slope)) + + elif isinstance(data, ACEtk.continuous.EnergyAngleDistributionData): + h5_group.attrs["type"] = "energy-angle-tabulated" + + if not data.interpolation_data.is_linear_linear: + print_error( + "Non-linearly-interpolated correlated-energy-angle is not supported" + ) + + NE = data.number_incident_energies + energy = _as_np(data.incident_energies) + h5_group.create_dataset("energy", data=energy).attrs["unit"] = "MeV" + + offset = np.zeros(NE, dtype=int) + energy_out = [] + pdf = [] + cosine_offset = [] + cosine = [] + cosine_pdf = [] + for i, distribution in enumerate(data.distributions): + offset[i] = len(pdf) + energy_out.extend(distribution.outgoing_energies) + pdf.extend(distribution.pdf) + + for inner_distribution in distribution.distributions: + cosine_offset.append(len(cosine_pdf)) + cosine.extend(inner_distribution.cosines) + cosine_pdf.extend(inner_distribution.pdf) + + h5_group.create_dataset("offset", data=offset) + h5_group.create_dataset("energy_out", data=np.array(energy_out)).attrs[ + "unit" + ] = "MeV" + h5_group.create_dataset("pdf", data=np.array(pdf)) + h5_group.create_dataset("cosine_offset", data=np.array(cosine_offset)) + h5_group.create_dataset("cosine", data=np.array(cosine)) + h5_group.create_dataset("cosine_pdf", data=np.array(cosine_pdf)) + + elif isinstance(data, ACEtk.continuous.NBodyPhaseSpaceDistribution): + h5_group.attrs["type"] = "N-body" + + if data.interpolation != 2: + print_error("Non-linearly-interpolable N-body energy distribution") + + h5_group.create_dataset("value", data=_as_np(data.values)).attrs["unit"] = "MeV" + h5_group.create_dataset("pdf", data=_as_np(data.pdf)) elif isinstance(data, ACEtk.continuous.SimpleMaxwellianFissionSpectrum): h5_group.attrs["type"] = "maxwellian" @@ -330,6 +557,46 @@ def load_energy_distribution(data, h5_group: h5py.Group): dataset.attrs["unit"] = "MeV" h5_group.create_dataset("pdf", data=data.pdf) + elif isinstance(data, ACEtk.continuous.MultiDistributionData): + h5_group.attrs["type"] = "multi" + + N = int(data.number_distributions) + h5_group.create_dataset("number_distributions", data=N) + + # Store probabilities for each sub-distribution + prob_group = h5_group.create_group("probabilities") + for i in range(N): + p = data.probability(i + 1) + g = prob_group.create_group(f"distribution-{i+1}") + + E = _as_np(p.energies) + P = _as_np(p.probabilities) + + g.create_dataset("energy", data=E).attrs["unit"] = "MeV" + g.create_dataset("probability", data=P) + + # Save interpolation metadata + b = _as_np(p.interpolation_data.boundaries).astype(int) + it = _as_np(p.interpolation_data.interpolants).astype(int) + g.create_dataset("interp_boundaries", data=b) + g.create_dataset("interp_interpolants", data=it) + + # Optional convenience: resample probability onto incident_grid if given + if incident_grid is not None and ( + not p.interpolation_data.is_linear_linear + ): + xg = np.asarray(incident_grid, dtype=float) + pres = np.array([_tab1_eval(x, E, P, b, it) for x in xg], dtype=float) + g.create_dataset("probability_on_incident_grid", data=pres) + + # Store each outgoing-energy distribution recursively + dist_group = h5_group.create_group("distributions") + for i in range(N): + sub = dist_group.create_group(f"distribution-{i+1}") + load_energy_distribution( + data.distribution(i + 1), sub, incident_grid=incident_grid + ) + else: print_error(f"Unsupported energy distribution: {data}")