{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Advanced McStas features: EXTEND and WHEN\n",
    "In this tutorial we will look at two advanced features in McStas, the EXTEND block and WHEN condition. Here we will use them to flag certain neutrons with EXTEND, and only record them in monitors when the flag is set using a WHEN condition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mcstasscript as ms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instrument = ms.McStas_instr(\"python_tutorial\", input_path=\"run_folder\",\n",
    "                             output_path=\"data_folder/mcstas_EXTEND_WHEN\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up an example McStas instrument\n",
    "First we set up an example instrument conisiting of a source, a guide and a position/divergence monitor. The guide is set up such that it only has mirrors on the left and right side, and absorbs neutrons if they hit the top or bottom. This is done to look at the horizontal behavior independently from the vertical, as this is easier to analyze."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "src = instrument.add_component(\"source\", \"Source_simple\")\n",
    "src.focus_xw = guide_opening_w = 0.05\n",
    "src.focus_yh = guide_opening_h = 0.06\n",
    "src.dist = source_guide_distance = 1.5\n",
    "\n",
    "wavelength = instrument.add_parameter(\"wavelength\", value=5.0, comment=\"Wavelength in [Ang]\")\n",
    "src.set_parameters(xwidth=0.02, yheight=0.02, flux=1E13, lambda0=wavelength, dlambda=\"0.001*wavelength\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "guide = instrument.add_component(\"guide\", \"Guide_gravity\")\n",
    "guide.set_AT([0, 0, source_guide_distance], RELATIVE=src)\n",
    "guide.set_parameters(w1=guide_opening_w, w2=guide_opening_w,\n",
    "                     h1=guide_opening_h, h2=guide_opening_h,\n",
    "                     l=15, G=-9.82, mleft=4.0, mright=4.0, mtop=0.0, mbottom=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "acceptance = instrument.add_component(\"acceptance\", \"DivPos_monitor\")\n",
    "acceptance.set_AT([0,0, guide.l + 0.1], RELATIVE=guide)\n",
    "acceptance.set_parameters(xwidth=0.08, yheight=0.05, nh=200, maxdiv_h=1.5, ndiv=200,\n",
    "                          filename='\"acceptance.dat\"', restore_neutron=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "scroll-output"
    ]
   },
   "outputs": [],
   "source": [
    "instrument.set_parameters(wavelength=2.8)\n",
    "instrument.settings(ncount=5E6)\n",
    "data = instrument.backengine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ms.make_sub_plot(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpreting the data\n",
    "Here we see an acceptance monitor, with position along the x-axis and divergence along the y-axis. The guide is under illuminated by the small source, so there are gaps in the acceptance diagram. We see the position and divergence of the beam consist of a large number of stripes, the ones with lowest divergence has the largest intensity.\n",
    "\n",
    "## Add an flag\n",
    "A flag is just a name for a variable that records some information on the neutron during the simulation, and can be used later to make a decision. Here we could check how many times the ray was reflected in the guide.\n",
    "\n",
    "We use an EXTEND block after a component to access variables internal to the component in the instrument scope. We declare a variable in the instrument scope called *n_reflections*. In the component scope, one can use the SCATTERED variable which contains the number of times the ray has encountered the SCATTER keyword within the component. Usually this is done when entering and leaving, and under each scattering / reflection, so the number of reflections is SCATTERED - 2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instrument.add_declare_var(\"int\", \"n_reflections\")\n",
    "guide.append_EXTEND(\"n_reflections = SCATTERED - 2;\")\n",
    "print(guide)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use the flag to limit what is recorded in a monitor\n",
    "A WHEN statement can be used to activate / deactivate a component when some condition is true / false. For example we could require 0 reflection in our guide. We add a few monitors similar to the original, with the only difference being WHEN statements requiring 0, 1 or 2 reflections in the guide for the component to be active. We use a for loop to add the similar components, only changing the component instance name, filename and WHEN statement between each."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reflection_numbers = [0, 1, 2]\n",
    "\n",
    "for reflections in reflection_numbers:\n",
    "    reflections_string = str(reflections)\n",
    "    \n",
    "    component_name = \"acceptance_\" + reflections_string\n",
    "    acceptance = instrument.add_component(component_name, \"DivPos_monitor\")\n",
    "    acceptance.filename = '\"acceptance_' + reflections_string + '.dat\"'\n",
    "    acceptance.set_WHEN(\"n_reflections == \" + reflections_string)\n",
    "    \n",
    "    acceptance.set_AT([0,0, guide.l + 0.1], RELATIVE=guide)\n",
    "    acceptance.set_parameters(xwidth=0.08, yheight=0.05, nh=200, maxdiv_h=1.5, ndiv=200, restore_neutron=1)\n",
    "    \n",
    "    print(acceptance)\n",
    "    print(\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The instrument diagram shows the use of *EXTEND* and *WHEN* keywords by decorating the corresponding component boxes with a heavier outline or dashed outline respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instrument.show_diagram()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running the simulation\n",
    "We now run the simulation with the new monitors to see how they differ from the original version."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "scroll-output"
    ]
   },
   "outputs": [],
   "source": [
    "instrument.set_parameters(wavelength=2.8)\n",
    "instrument.settings(ncount=5E6)\n",
    "instrument.show_settings()\n",
    "\n",
    "new_data = instrument.backengine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ms.make_sub_plot(new_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpretation of the data\n",
    "The original monitor is unchanged as it was not modified. On the monitors with different numbers of reflections, we see the middle line correspond to zero reflections, the two lines around those are for one reflection and so forth. This explains why the lines further from the center has lower intensity, as they underwent more reflections while also having a larger angle of incidence on the mirrors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The McStas instrument file\n",
    "We here show the generated McStas instrument file in order to clarify how this would be accomplished without the McStasScript API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instrument.show_instrument_file()"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}