diff --git a/k_eff_search_with_tally_derivatives.ipynb b/k_eff_search_with_tally_derivatives.ipynb new file mode 100644 index 0000000..5a93481 --- /dev/null +++ b/k_eff_search_with_tally_derivatives.ipynb @@ -0,0 +1,483 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "fff957fb", + "metadata": {}, + "source": [ + "# Using tally derivatives for keff search: An alternative to search_for_keff (openmc python API).\n", + "\n", + "This notebook demonstrates a gradient-based approach using OpenMC tally derivatives to find a critical boron concentration (ppm) and compares it with the built-in `openmc.search_for_keff` method." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "43291742", + "metadata": {}, + "outputs": [], + "source": [ + "# Core imports and configuration\n", + "import os\n", + "import openmc\n", + "import numpy as np\n", + "import warnings\n", + "\n", + "\n", + "# Suppress FutureWarnings for cleaner output\n", + "warnings.filterwarnings('ignore', category=FutureWarning)\n", + "\n", + "# Physical/constants used in chain-rule conversions\n", + "N_A = 6.02214076e23 # Avogadro's number (atoms/mol)\n", + "A_B_nat = 10.81 # g/mol approximate atomic mass for natural boron\n", + "\n", + "print(\"PATH:\", os.environ.get('PATH'))\n", + "print(\"OPENMC_CROSS_SECTIONS:\", os.environ.get('OPENMC_CROSS_SECTIONS'))\n", + "os.environ['PATH'] = '/workspaces/openmc/build/bin/:' + os.environ['PATH']\n", + "os.environ['OPENMC_CROSS_SECTIONS'] = '/home/codespace/nndc_hdf5/cross_sections.xml'\n" + ] + }, + { + "cell_type": "markdown", + "id": "a4ac511c", + "metadata": {}, + "source": [ + "**Model builder**: This cell contains the `build_model(ppm_Boron)` function that constructs materials, geometry, and settings for the simple pin-cell model used in the experiments. Keep this function as-is to ensure reproducibility." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "88e65308", + "metadata": {}, + "outputs": [], + "source": [ + "# ===============================================================\n", + "# Model builder\n", + "# ===============================================================\n", + "def build_model(ppm_Boron):\n", + " # Create the pin materials\n", + " fuel = openmc.Material(name='1.6% Fuel', material_id=1)\n", + " fuel.set_density('g/cm3', 10.31341)\n", + " fuel.add_element('U', 1., enrichment=1.6)\n", + " fuel.add_element('O', 2.)\n", + "\n", + " zircaloy = openmc.Material(name='Zircaloy', material_id=2)\n", + " zircaloy.set_density('g/cm3', 6.55)\n", + " zircaloy.add_element('Zr', 1.)\n", + "\n", + " water = openmc.Material(name='Borated Water', material_id=3)\n", + " water.set_density('g/cm3', 0.741)\n", + " water.add_element('H', 2.)\n", + " water.add_element('O', 1.)\n", + " water.add_element('B', ppm_Boron * 1e-6)\n", + "\n", + " materials = openmc.Materials([fuel, zircaloy, water])\n", + "\n", + " # Geometry\n", + " fuel_outer_radius = openmc.ZCylinder(r=0.39218)\n", + " clad_outer_radius = openmc.ZCylinder(r=0.45720)\n", + "\n", + " min_x = openmc.XPlane(x0=-0.63, boundary_type='reflective')\n", + " max_x = openmc.XPlane(x0=+0.63, boundary_type='reflective')\n", + " min_y = openmc.YPlane(y0=-0.63, boundary_type='reflective')\n", + " max_y = openmc.YPlane(y0=+0.63, boundary_type='reflective')\n", + "\n", + " fuel_cell = openmc.Cell(name='1.6% Fuel')\n", + " fuel_cell.fill = fuel\n", + " fuel_cell.region = -fuel_outer_radius\n", + "\n", + " clad_cell = openmc.Cell(name='1.6% Clad')\n", + " clad_cell.fill = zircaloy\n", + " clad_cell.region = +fuel_outer_radius & -clad_outer_radius\n", + "\n", + " moderator_cell = openmc.Cell(name='1.6% Moderator')\n", + " moderator_cell.fill = water\n", + " moderator_cell.region = +clad_outer_radius & (+min_x & -max_x & +min_y & -max_y)\n", + "\n", + " root_universe = openmc.Universe(name='root universe', universe_id=0)\n", + " root_universe.add_cells([fuel_cell, clad_cell, moderator_cell])\n", + "\n", + " geometry = openmc.Geometry(root_universe)\n", + "\n", + " # Settings\n", + " settings = openmc.Settings()\n", + " settings.batches = 300\n", + " settings.inactive = 20\n", + " settings.particles = 1000\n", + " settings.run_mode = 'eigenvalue'\n", + " settings.verbosity=1\n", + "\n", + " bounds = [-0.63, -0.63, -10, 0.63, 0.63, 10.]\n", + " uniform_dist = openmc.stats.Box(bounds[:3], bounds[3:], only_fissionable=True)\n", + " settings.source = openmc.Source(space=uniform_dist)\n", + "\n", + " model = openmc.model.Model(geometry, materials, settings)\n", + " return model" + ] + }, + { + "cell_type": "markdown", + "id": "4137128e", + "metadata": {}, + "source": [ + "**Helpers**: utility functions for extracting cell IDs used by a material and for running OpenMC with derivative tallies." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2e903305", + "metadata": {}, + "outputs": [], + "source": [ + "# helper to get the material id\n", + "def find_cells_using_material(geometry, material):\n", + " \"Return list of cell ids in `geometry` filled with `material`.\"\n", + " return [c.id for c in geometry.get_all_cells().values() if c.fill is material]\n" + ] + }, + { + "cell_type": "markdown", + "id": "ac69e986", + "metadata": {}, + "source": [ + "**Running OpenMC with derivative tallies**: `run_with_gradient` runs an OpenMC simulation for a given boron ppm, attaches derivative tallies for the boron isotopes, and returns k-eff plus the required tallies to compute dk/dppm. Keep this implementation intact so gradient calculations remain consistent with the original notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "703e0483", + "metadata": {}, + "outputs": [], + "source": [ + "def run_with_gradient(ppm_B, target_batches=50, water_material_id=3, boron_nuclides=('B10', 'B11')):\n", + " \"\"\"Run OpenMC and compute k-effective with gradient information\"\"\"\n", + " # Clean up previous files\n", + " for f in ['summary.h5', f'statepoint.{target_batches}.h5', 'tallies.out']:\n", + " if os.path.exists(f):\n", + " os.remove(f)\n", + "\n", + " # Build model\n", + " model = build_model(ppm_B)\n", + "\n", + " # Auto-detect moderator cells\n", + " water = model.materials[water_material_id - 1] # Materials are 0-indexed\n", + " moderator_cell_ids = find_cells_using_material(model.geometry, water)\n", + " moderator_filter = openmc.CellFilter(moderator_cell_ids)\n", + "\n", + " # Base tallies\n", + " tF_base = openmc.Tally(name='FissionBase')\n", + " tF_base.scores = ['nu-fission']\n", + " tA_base = openmc.Tally(name='AbsorptionBase')\n", + " tA_base.scores = ['absorption']\n", + "\n", + " # Derivative tallies\n", + " deriv_tallies = []\n", + " for nuc in boron_nuclides:\n", + " deriv = openmc.TallyDerivative(\n", + " variable='nuclide_density',\n", + " material=water_material_id,\n", + " nuclide=nuc\n", + " )\n", + "\n", + " tf = openmc.Tally(name=f'Fission_deriv_{nuc}')\n", + " tf.scores = ['nu-fission']\n", + " tf.derivative = deriv\n", + " tf.filters = [moderator_filter]\n", + "\n", + " ta = openmc.Tally(name=f'Absorp_deriv_{nuc}')\n", + " ta.scores = ['absorption']\n", + " ta.derivative = deriv\n", + " ta.filters = [moderator_filter]\n", + "\n", + " deriv_tallies += [tf, ta]\n", + "\n", + " model.tallies = openmc.Tallies([tF_base, tA_base] + deriv_tallies)\n", + " model.settings.batches = target_batches\n", + " model.settings.inactive = max(1, int(target_batches * 0.1))\n", + "\n", + " # Run simulation\n", + " model.run()\n", + " sp = openmc.StatePoint(f\"statepoint.{target_batches}.h5\")\n", + "\n", + " # Get results\n", + " k_eff = sp.keff.nominal_value\n", + "\n", + " # Base tallies\n", + " fission_tally = sp.get_tally(name='FissionBase')\n", + " absorption_tally = sp.get_tally(name='AbsorptionBase')\n", + " F_base = float(np.sum(fission_tally.mean))\n", + " A_base = float(np.sum(absorption_tally.mean))\n", + "\n", + " # Derivative tallies\n", + " dF_dN_total = 0.0\n", + " dA_dN_total = 0.0\n", + "\n", + " for nuc in boron_nuclides:\n", + " fission_deriv = sp.get_tally(name=f'Fission_deriv_{nuc}')\n", + " absorption_deriv = sp.get_tally(name=f'Absorp_deriv_{nuc}')\n", + "\n", + " if fission_deriv:\n", + " dF_dN_total += float(np.sum(fission_deriv.mean))\n", + " if absorption_deriv:\n", + " dA_dN_total += float(np.sum(absorption_deriv.mean))\n", + "\n", + " return k_eff, F_base, A_base, dF_dN_total, dA_dN_total, water.density" + ] + }, + { + "cell_type": "markdown", + "id": "1dc12f6f", + "metadata": {}, + "source": [ + "**Gradient-based optimizer**: the gradient descent routine that uses the analytical derivative tallies to propose ppm updates. The original adaptive logic is preserved; we add an experiments cell later to vary tuning parameters.\n", + "\n", + "#### Theory, derivation and scaling\n", + "\n", + "This optimizer uses derivative tallies to estimate how small changes in boron concentration (ppm) affect the reactor multiplication factor k_eff. The notebook treats k approximately as the ratio of a fission production tally F to an absorption tally A, i.e. k ≈ F / A. Both F and A depend on the boron number density N (atoms/cm³).\n", + "\n", + "From calculus (quotient rule) we get the derivative of k with respect to N:\n", + "\n", + "- dk/dN = (A * dF_dN - F * dA_dN) / A^2\n", + "\n", + "This expression follows directly from d(f/g)/dN = (g f' - f g') / g^2 (standard calculus — see any calculus text).\n", + "\n", + "To convert from mass-part-per-million (ppm) to number density we use the mass fraction and Avogadro's number. If ppm is a mass-part-per-million, the boron mass fraction is ppm × 1e-6 and the boron number density is\n", + "\n", + "- N = (rho_water * mass_fraction) × (N_A / A_B) \n", + "so\n", + "- dN/dppm = 1e-6 × rho_water × N_A / A_B\n", + "\n", + "Finally, combine with the chain rule: dk/dppm = dk/dN × dN/dppm.\n", + "\n", + "\n", + "\n", + "#### Practical considerations and recommendations\n", + "\n", + "- Large dk/dppm magnitudes are expected numerically because ppm is a very small mass fraction but corresponds to a large change in atom counts when converted to atoms/cm³. Expect amplification by ~1e16–1e17 from the conversion alone.\n", + "- When using these derivatives in an optimizer you should scale or clip updates to keep steps physically reasonable (e.g., clamp per-iteration ppm changes, perform line search/backtracking, optimize over log(ppm) rather than linear ppm, or normalize the gradient to a target step magnitude).\n", + "- Also account for stochastic noise: derivative tallies from Monte Carlo will have statistical uncertainty. Increase batches/particles or use averaging/adjoint methods for more robust gradients." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "36639b85", + "metadata": {}, + "outputs": [], + "source": [ + "def gradient_based_search(ppm_start, ppm_range, k_target, max_iter, tol=1e-2, learning_rate=1e-17):\n", + " \"\"\"Gradient-based optimization using analytical derivatives\"\"\"\n", + " ppm = float(ppm_start)\n", + " history = []\n", + "\n", + " print(\"GRADIENT-BASED OPTIMIZATION\")\n", + " print(f\"Initial: {ppm:.1f} ppm, Target: k = {k_target}\")\n", + " print(\"Iter | ppm | k_eff | Error | Gradient | Step | Learning Rate\")\n", + " print(\"-\" * 85)\n", + "\n", + " for it in range(max_iter):\n", + " k, F, A, dF_dN, dA_dN, rho_water = run_with_gradient(ppm)\n", + " err = k - k_target\n", + " history.append((ppm, k, err, dF_dN, dA_dN, learning_rate))\n", + "\n", + " # Calculate gradient using chain rule\n", + " dk_dN = (A * dF_dN - F * dA_dN) / (A * A)\n", + " dN_dppm = 1e-6 * rho_water * N_A / A_B_nat\n", + " dk_dppm = dk_dN * dN_dppm \n", + "\n", + " # Gradient descent step with momentum-like behavior for small gradients\n", + " if abs(dk_dppm) < 1e-10: # Very small gradient\n", + " # Use a conservative fixed step in the right direction\n", + " step = -100 if err > 0 else 100\n", + " else:\n", + " step = -learning_rate * err * dk_dppm\n", + "\n", + " # Additional adaptive scaling based on error magnitude\n", + " error_magnitude = abs(err)\n", + " if error_magnitude > 0.1:\n", + " step *= 1.5\n", + " elif error_magnitude < 0.01:\n", + " step *= 0.7\n", + "\n", + " ppm_new = ppm + step\n", + " print(\"NEW PPM SUGGESTION: \", ppm_new)\n", + " print(\"ERROR WAS: \", err) \n", + " # Apply bounds\n", + " ppm_new = max(ppm_range[0], min(ppm_new, ppm_range[1]))\n", + "\n", + "\n", + " print(f\"{it+1:3d} | {ppm:7.1f} | {k:9.6f} | {err:7.4f} | {dk_dppm:10.2e} | {step:7.1f} | {learning_rate:12.2e}\")\n", + "\n", + " if abs(err) < tol:\n", + " print(f\"✓ CONVERGED in {it+1} iterations\")\n", + " return ppm, history\n", + "\n", + " ppm = ppm_new\n", + "\n", + " print(f\"Reached maximum iterations ({max_iter})\")\n", + " return ppm, history" + ] + }, + { + "cell_type": "markdown", + "id": "dcab00f6", + "metadata": {}, + "source": [ + "**OpenMC built-in keff search**: wrapper around `openmc.search_for_keff` used for baseline comparisons." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "afebca98", + "metadata": {}, + "outputs": [], + "source": [ + "def builtin_keff_search(k_target, ppm_start, ppm_range, max_iter):\n", + " \"\"\"Call `openmc.search_for_keff` with a small bracket and return results.\"\"\"\n", + " print(\"\\n===== OPENMC BUILTIN KEFF SEARCH =====\\n\")\n", + "\n", + " crit_ppm, guesses, keffs = openmc.search_for_keff(\n", + " build_model,\n", + " initial_guess = ppm_start,\n", + " target=k_target,\n", + " bracket=ppm_range,\n", + " tol=1e-2,\n", + " print_iterations=True,\n", + " run_args={'output': False}\n", + " )\n", + "\n", + " print(\"\\nCritical Boron Concentration: {:4.0f} ppm\".format(crit_ppm))\n", + " return crit_ppm, guesses, keffs" + ] + }, + { + "cell_type": "markdown", + "id": "a6fb4365", + "metadata": {}, + "source": [ + "**Comparison function**: run both methods and summarize results. This function keeps the previous comparison logic intact." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a72b897a", + "metadata": {}, + "outputs": [], + "source": [ + "def compare_optimization_methods(ppm_start, k_target, ppm_range, max_iter):\n", + " \"\"\"Compare gradient-based vs built-in search and summarize results.\"\"\"\n", + " print(\"=\" * 80)\n", + " print(\"COMPARING OPTIMIZATION METHODS FOR BORON CONCENTRATION SEARCH\")\n", + " print(f\"Target k_eff: {k_target}, Initial guess: {ppm_start} ppm\")\n", + " print(\"=\" * 80)\n", + " \n", + " # Method 1: OpenMC function (gradient-free)\n", + " print(\"\\n=== Running OpenMC built-in keff search ===\")\n", + " builtin_ppm, guesses, keffs = builtin_keff_search(k_target, ppm_start, ppm_range, max_iter)\n", + " # Convert built-in search logs to unified history format (approximate)\n", + " builtin_history = [(g, k, k - 1.0, 0, 0) for g, k in zip(guesses, keffs)]\n", + " \n", + " # Method 2: Gradient-based (analytical derivatives)\n", + " grad_ppm, grad_history = gradient_based_search(ppm_start, ppm_range, k_target, max_iter)\n", + "\n", + " methods = [\n", + " (\"Analytical Gradient\", grad_ppm, grad_history),\n", + " (\"OpenMC Built-in\", builtin_ppm, builtin_history),\n", + " ]\n", + "\n", + " # Results comparison\n", + " print(\"\\n\" + \"=\" * 80)\n", + " print(\"FINAL RESULTS COMPARISON\")\n", + " print(\"=\" * 80)\n", + "\n", + " best_method = None\n", + " best_error = float('inf')\n", + "\n", + " for name, ppm, history in methods:\n", + " if history:\n", + " final_k = history[-1][1]\n", + " final_err = abs(history[-1][2])\n", + " iterations = len(history)\n", + " print(f\"\\n{name}:\")\n", + " print(f\" Final ppm: {ppm:.1f}\")\n", + " print(f\" Final k_eff: {final_k:.6f}\")\n", + " print(f\" Final error: {final_err:.6f}\")\n", + " print(f\" Iterations: {iterations}\")\n", + " if final_err < best_error:\n", + " best_error = final_err\n", + " best_method = name\n", + "\n", + " if best_method:\n", + " print(f\"\\n★ BEST METHOD: {best_method} (error = {best_error:.6f})\")\n", + "\n", + " # Convergence speed analysis\n", + " print(f\"\\nCONVERGENCE SPEED ANALYSIS:\")\n", + " tolerance_levels = [0.05, 0.02, 0.01] # 5%, 2%, 1% tolerance\n", + " for name, ppm, history in methods:\n", + " if history:\n", + " print(f\"\\n{name}:\")\n", + " for tol_level in tolerance_levels:\n", + " iterations_to_tolerance = None\n", + " for i, (_, k, err, *_) in enumerate(history):\n", + " if abs(err) < tol_level:\n", + " iterations_to_tolerance = i + 1\n", + " break\n", + " if iterations_to_tolerance:\n", + " print(f\" Reached {tol_level*100:.0f}% tolerance in {iterations_to_tolerance} iterations\")\n", + " else:\n", + " print(f\" Did not reach {tol_level*100:.0f}% tolerance\")\n", + "\n", + " return methods" + ] + }, + { + "cell_type": "markdown", + "id": "d35c178a", + "metadata": {}, + "source": [ + "## Lets run the iterative gradient search version and compare it with the gradient-free (python api) k_eff search approach" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "756e671c", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "ppm_range = [0.0005, 20000]\n", + "ppm_start = ppm_range[0]\n", + "k_target = 0.95\n", + "max_iter = 50\n", + "compare_optimization_methods(ppm_start, k_target, ppm_range, max_iter)\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.1" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}