{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2e8a7f97-24b4-455e-b768-57cdee820012",
   "metadata": {},
   "source": [
    "# Create perturbed FY files (only variance, all energies)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "06dad8ae-478b-44fe-a0cc-8cd37bd4acc3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-03T21:51:02.354008Z",
     "iopub.status.busy": "2024-03-03T21:51:02.353529Z",
     "iopub.status.idle": "2024-03-03T21:51:02.969034Z",
     "shell.execute_reply": "2024-03-03T21:51:02.968511Z"
    }
   },
   "outputs": [],
   "source": [
    "import sandy\n",
    "import pandas as pd\n",
    "from os.path import join\n",
    "import numpy as np\n",
    "import random, sys"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db66e7a6-e7a9-46a2-abea-6cfc699122b3",
   "metadata": {},
   "source": [
    "## Extract FYs and covariance data for U235 fission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3096afe7-8921-4dfe-94b4-1e3443954424",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-03T21:51:02.971694Z",
     "iopub.status.busy": "2024-03-03T21:51:02.971278Z",
     "iopub.status.idle": "2024-03-03T21:51:03.673060Z",
     "shell.execute_reply": "2024-03-03T21:51:03.672423Z"
    }
   },
   "outputs": [],
   "source": [
    "za = 92235\n",
    "tape = sandy.get_endf6_file(\"jeff_33\", \"nfpy\", za * 10)\n",
    "nfpy = sandy.Fy.from_endf6(tape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "891967d7-eb66-424f-a8a4-14f625f133cf",
   "metadata": {},
   "source": [
    "## Generate perturbation coefficients for all energies and write them to file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ec1dd824-00f2-4b21-94f4-38728e586b12",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-03T21:51:03.675686Z",
     "iopub.status.busy": "2024-03-03T21:51:03.675440Z",
     "iopub.status.idle": "2024-03-03T21:51:03.678695Z",
     "shell.execute_reply": "2024-03-03T21:51:03.678124Z"
    }
   },
   "outputs": [],
   "source": [
    "energies = nfpy.data.E.unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7b45981f-8bf9-4c4b-bbb5-5f3e8ed4098a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-03T21:51:03.680740Z",
     "iopub.status.busy": "2024-03-03T21:51:03.680405Z",
     "iopub.status.idle": "2024-03-03T21:51:03.956277Z",
     "shell.execute_reply": "2024-03-03T21:51:03.955683Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:  Large condition number of covariance matrix: 1.36e+11\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:  Large condition number of covariance matrix: 1.38e+11\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sampling IFY for energy 2.530e-02 eV...\n",
      "sampling IFY for energy 4.000e+05 eV...\n",
      "sampling IFY for energy 1.400e+07 eV...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:  Large condition number of covariance matrix: 1.28e+11\n"
     ]
    }
   ],
   "source": [
    "nsmp = 10 # sample size\n",
    "smp = {}\n",
    "for e in energies:\n",
    "    idx = nfpy.data.query(f\"E=={e} & MT==454\").index\n",
    "    ify = nfpy.data.loc[idx]\n",
    "    cov = sandy.CategoryCov(np.diag(ify.DFY**2), index=ify.ZAP, columns=ify.ZAP)  # Diagonal covariance matrix\n",
    "    seed = random.randrange(2**32 - 1)    # create a seed\n",
    "    print(f\"sampling IFY for energy {e:.3e} eV...\")\n",
    "    smp[e] = cov.sampling(nsmp, seed=seed)   # need to change the seed for the different energies    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1d6bb1ef-e27b-4e2f-91b5-29156a293860",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-03T21:51:03.959128Z",
     "iopub.status.busy": "2024-03-03T21:51:03.958685Z",
     "iopub.status.idle": "2024-03-03T21:51:05.012513Z",
     "shell.execute_reply": "2024-03-03T21:51:05.011982Z"
    }
   },
   "outputs": [],
   "source": [
    "with pd.ExcelWriter(f'PERT_{za}_MF8_MT454.xlsx') as writer:\n",
    "    for e, s in smp.items():\n",
    "        s.data.to_excel(writer, sheet_name=f'{e:.3e}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9d980cd-c181-465a-a499-acfaf0777027",
   "metadata": {},
   "source": [
    "## Read coefficients from perturbation file and generate random FY ENDF-6 files"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57c30635-895e-4ded-8b0d-ae9933f6cbf6",
   "metadata": {},
   "source": [
    "Skip the part above if you already have the file of perturbations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e2eec6f2-8129-4395-a921-fdb13ca043ed",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-03T21:51:05.015149Z",
     "iopub.status.busy": "2024-03-03T21:51:05.014805Z",
     "iopub.status.idle": "2024-03-03T21:51:05.383741Z",
     "shell.execute_reply": "2024-03-03T21:51:05.383087Z"
    }
   },
   "outputs": [],
   "source": [
    "smp = pd.read_excel(f'PERT_{za}_MF8_MT454.xlsx', sheet_name=None, index_col=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1e104351-9068-4d8b-adbb-9fefa18407b9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-03T21:51:05.386153Z",
     "iopub.status.busy": "2024-03-03T21:51:05.385947Z",
     "iopub.status.idle": "2024-03-03T21:51:05.911265Z",
     "shell.execute_reply": "2024-03-03T21:51:05.910615Z"
    }
   },
   "outputs": [],
   "source": [
    "za = 92235\n",
    "tape = sandy.get_endf6_file(\"jeff_33\", \"nfpy\", za * 10)\n",
    "nfpy = sandy.Fy.from_endf6(tape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ca0cd0b6-e679-48e9-87e9-dddf139a8a08",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-03T21:51:05.913916Z",
     "iopub.status.busy": "2024-03-03T21:51:05.913565Z",
     "iopub.status.idle": "2024-03-03T21:51:05.916281Z",
     "shell.execute_reply": "2024-03-03T21:51:05.915700Z"
    }
   },
   "outputs": [],
   "source": [
    "### run only if you want consistent CFYs\n",
    "# tape_rdd = sandy.get_endf6_file(\"jeff_33\", \"decay\", \"all\")\n",
    "# rdd = sandy.DecayData.from_endf6(tape_rdd)  # this can take a while"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "181943ff-d4b9-4ba1-9fa3-94c67092f68c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-03T21:51:05.918510Z",
     "iopub.status.busy": "2024-03-03T21:51:05.918154Z",
     "iopub.status.idle": "2024-03-03T21:51:11.118846Z",
     "shell.execute_reply": "2024-03-03T21:51:11.118285Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "writing file 'u235_fy_0.jeff33'...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "writing file 'u235_fy_1.jeff33'...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "writing file 'u235_fy_2.jeff33'...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "writing file 'u235_fy_3.jeff33'...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "writing file 'u235_fy_4.jeff33'...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "writing file 'u235_fy_5.jeff33'...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "writing file 'u235_fy_6.jeff33'...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "writing file 'u235_fy_7.jeff33'...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "writing file 'u235_fy_8.jeff33'...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "writing file 'u235_fy_9.jeff33'...\n"
     ]
    }
   ],
   "source": [
    "smp_min = 0   # write ENDF-6 file only in the sample range [smp_min, smp_max]\n",
    "smp_max = 9\n",
    "file_template = \"u235_fy_{}.jeff33\"\n",
    "for ismp in range(smp_min, smp_max+1):\n",
    "    file = file_template.format(ismp)\n",
    "    f = sandy.Fy(nfpy.data.copy())\n",
    "    for e, s in smp.items():\n",
    "        idx_ify = nfpy.data.query(f\"E=={float(e)} & MT==454\").index\n",
    "        idx_cfy = nfpy.data.query(f\"E=={float(e)} & MT==459\").index\n",
    "        f.data.loc[idx_ify, \"DFY\"] = f.data.loc[idx_ify, \"FY\"]            # just for me, i copy the original IFYs where uncertainties should be, so i can compare them to the perturbed ones (anyways I don't use uncertainties)\n",
    "        f.data.loc[idx_cfy, \"DFY\"] = f.data.loc[idx_cfy, \"FY\"]            # same but for CFYs\n",
    "        f.data.loc[idx_ify, \"FY\"] *= s[ismp].values                # IMPORTANT, this does not update the CFYs, which in random ENDF-6 file are inconsistent with the perturbed IFYs\n",
    "        #f = f.apply_qmatrix(922350, e, rdd, keep_fy_index=True)      # Run this if you want to update the CFYs (slower), or else comment it out\n",
    "    print(f\"writing file '{file}'...\")\n",
    "    f.to_endf6(tape).to_file(file)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sandy-devel",
   "language": "python",
   "name": "sandy-devel"
  },
  "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}