{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Simulation and Advanced Tricks\n",
    "\n",
    "This notebook will demonstrate how one can use the modular structure of fuse to change parts of the simulation by exchanging or adding new plugins. To get the most out of this notebook, you should start with an empty `out_dir` when setting up the simulation context. Additionaly the notebook will show how a map can be replaced with a dummy map to disable certain parts of the simulation.\n",
    "\n",
    "## Imports & Simulation Context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import fuse\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st = fuse.context.xenonnt_fuse_full_chain_simulation(\n",
    "    output_folder=\"./fuse_data\",\n",
    "    simulation_config=\"sr0_dev\",\n",
    ")\n",
    "\n",
    "st.set_config(\n",
    "    {\n",
    "        \"path\": \"/project2/lgrandi/xenonnt/simulations/testing\",\n",
    "        \"file_name\": \"pmt_neutrons_100.root\",\n",
    "        \"entry_stop\": 10,\n",
    "        \"debug\": True,\n",
    "    }\n",
    ")\n",
    "\n",
    "run_number = \"00000\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the default simulation\n",
    "First we will run the default microphysics simulation. As we set the `debug` config to `True` fuse will give use some more information during the simulation. For example, each plugin will print its version number and the used seed for random number generation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st.make(run_number, \"microphysics_summary\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using a different plugin\n",
    "\n",
    "First lets see how we can exchange a fuse plugin with a plugin that is included in fuse but not used by default. One example is the `BBFYields` plugin that can be used instead of `NestYields`. To add a plugin to our simulation context we can use the `st.register` function provided by strax. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st.register(fuse.plugins.micro_physics.yields.BBFYields)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that the new plugin is registered we can run the simulation again. Can you spot the difference in the debug output?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st.make(run_number, \"microphysics_summary\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "fuse should now tell you:\n",
    "```\n",
    "DEBUG:fuse.micro_physics.yields:Running BBFYields version 0.1.0 in debug mode\n",
    "```\n",
    "This way you can easily see that the new plugin is used. You do not need to worry that you could mix up data produced with different plugins as strax recognizes the changed simulation context and will rerun the needed simulation steps. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a new plugin\n",
    "\n",
    "Now that we know how to exchange a plugin in the context we can build a new plugin and use it in the simulation. For our example we will replace the `ElectronDrift` plugin with a new plugin that does the simulation a \"creative\" way. You can find the plugin with some inline comments below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import strax\n",
    "import straxen\n",
    "import numpy as np\n",
    "import logging\n",
    "\n",
    "logging.basicConfig(handlers=[logging.StreamHandler()])\n",
    "log = logging.getLogger(\"fuse.detector_physics.electron_drift\")\n",
    "\n",
    "from fuse.common import FUSE_PLUGIN_TIMEOUT\n",
    "\n",
    "\n",
    "# A fuse plugin is a python class that inherits from strax.Plugin\n",
    "# As naming convention we use CamelCase for the class name\n",
    "class CosineElectronDrift(strax.Plugin):\n",
    "    # Each plugin has a version number\n",
    "    # If the version number changes, fuse will know that it need to re-simulate the data\n",
    "    __version__ = \"0.1.0\"\n",
    "\n",
    "    # You need to tell fuse and strax what the plugin needs as input\n",
    "    # In this case we need the microphysics_summary\n",
    "    # If you need more than one input, you can use a tuple\n",
    "    depends_on = \"microphysics_summary\"\n",
    "\n",
    "    # You need to tell fuse and strax what the plugin provides as output\n",
    "    # In this case we provide drifted_electrons\n",
    "    # You can later use st.make(run_number, \"drifted_electrons\") to run the simulation\n",
    "    provides = \"drifted_electrons\"\n",
    "\n",
    "    # You need to tell fuse and strax what the data looks like\n",
    "    # Data of the same data_kind can be combined via \"horizontal\" concatenation and need\n",
    "    # to have the same output length.\n",
    "    data_kind = \"interactions_in_roi\"\n",
    "\n",
    "    # You also need to tell strax what columns the data has\n",
    "    # A column needs a name and a numpy data type.\n",
    "    # For this example we will not deal with the other colums usually present in drifted_electrons\n",
    "    # This can lead to problems later on in the simulation but is fine for this example\n",
    "    dtype = [\n",
    "        (\"n_electron_interface\", np.int64),\n",
    "    ]\n",
    "    dtype = dtype + strax.time_fields\n",
    "\n",
    "    # We need to disable automatic rechunking for fuse plugins\n",
    "    # As fuse is going from \"leightweigt\" data to \"heavy\" data,\n",
    "    # automatic rechunking can lead to problems in later plugins\n",
    "    rechunk_on_save = False\n",
    "\n",
    "    # We need to specify when we want to save the data\n",
    "    save_when = strax.SaveWhen.TARGET\n",
    "\n",
    "    # strax uses a rather short timeout, lets increase it as\n",
    "    # some of the fuse simulation steps can take a while\n",
    "    input_timeout = FUSE_PLUGIN_TIMEOUT\n",
    "\n",
    "    # We need to tell strax what config options the plugin needs\n",
    "    # We will use the great URLConfigs that are a part of straxen\n",
    "    debug = straxen.URLConfig(\n",
    "        default=False,\n",
    "        type=bool,\n",
    "        track=False,\n",
    "        help=\"Show debug informations\",\n",
    "    )\n",
    "\n",
    "    deterministic_seed = straxen.URLConfig(\n",
    "        default=True,\n",
    "        type=bool,\n",
    "        help=\"Set the random seed from lineage and run_id, or pull the seed from the OS.\",\n",
    "    )\n",
    "\n",
    "    # For our example we will need the tpc_length\n",
    "    tpc_length = straxen.URLConfig(\n",
    "        type=(int, float),\n",
    "        help=\"tpc_length\",\n",
    "    )\n",
    "\n",
    "    # And a new variable we will call n_periods\n",
    "    n_periods = straxen.URLConfig(\n",
    "        type=(int, float),\n",
    "        help=\"n_periods\",\n",
    "    )\n",
    "\n",
    "    # If you want to prepare something before we start to run the compute method\n",
    "    # you can put it into the setup method. The setup method is called once while the\n",
    "    # compute method is called independently for each chunk\n",
    "    def setup(self):\n",
    "        # Lets convert the tpc_length and n_periods into a scaling factor\n",
    "        self.scaling_factor = self.tpc_length / self.n_periods\n",
    "\n",
    "        # All plugins can report problmes or debug information via the logging feature\n",
    "        # You can set the log level via the debug config option.\n",
    "        # WARNING messages are always shown whild DEBUG messages are only shown if debug is True\n",
    "        if self.debug:\n",
    "            log.setLevel(\"DEBUG\")\n",
    "            log.debug(f\"Running ElectronDrift version {self.__version__} in debug mode\")\n",
    "        else:\n",
    "            log.setLevel(\"WARNING\")\n",
    "\n",
    "        # Many plugins need to generate random numbers for simulation the corresponding physics process\n",
    "        # In fuse we want to make sure that the simulation is reproducible.\n",
    "        # Therefore we have the default setting of deterministic_seed = True\n",
    "        # In this case the random seed is generated from the run_id and the lineage\n",
    "        # The lineage includes all plugins and their verions that are connected to the input of the\n",
    "        # current plugin as well as all tracked config options and the strax version.\n",
    "        # The run_id is a user input. More on the deterministic seed can be found in\n",
    "        # a dedicated notebook.\n",
    "        # Please make sure that you use the random number generator self.rng when you need random numbers\n",
    "        # later in the plugin.\n",
    "        if self.deterministic_seed:\n",
    "            hash_string = strax.deterministic_hash((self.run_id, self.lineage))\n",
    "            seed = int(hash_string.encode().hex(), 16)\n",
    "            self.rng = np.random.default_rng(seed=seed)\n",
    "            log.debug(f\"Generating random numbers from seed {seed}\")\n",
    "        else:\n",
    "            self.rng = np.random.default_rng()\n",
    "            log.debug(f\"Generating random numbers with seed pulled from OS\")\n",
    "\n",
    "    # The compute method is the heart of the plugin. It is executed for each chunk of input data and\n",
    "    # must produce data in the format specified in the self.dtype variable.\n",
    "    def compute(self, interactions_in_roi):\n",
    "        # For your new plugin you would put your new simulation code here. In this example\n",
    "        # We will do some stupid calculation of the drift time.\n",
    "\n",
    "        # Make sure you only apply the plugin to interactions that have some electrons\n",
    "        # Adapt this to match your plugin-input\n",
    "        mask = interactions_in_roi[\"electrons\"] > 0\n",
    "\n",
    "        # Make sure your plugin can handle empty inputs\n",
    "        if len(interactions_in_roi[mask]) == 0:\n",
    "            return np.zeros(0, self.dtype)\n",
    "\n",
    "        # Build the output array with the correct length and dtype\n",
    "        result = np.zeros(len(interactions_in_roi), dtype=self.dtype)\n",
    "        # We do not want to change the timing of the interactions, so we just take them from the input\n",
    "        result[\"time\"] = interactions_in_roi[\"time\"]\n",
    "        result[\"endtime\"] = interactions_in_roi[\"endtime\"]\n",
    "\n",
    "        # Lets do some stupid calculation of the drifted electrons using a cosine function\n",
    "        n_electron_interface = (\n",
    "            interactions_in_roi[mask][\"electrons\"]\n",
    "            * np.cos(interactions_in_roi[mask][\"z\"] / self.scaling_factor * np.pi) ** 2\n",
    "        )\n",
    "        # Lets add some noise and make sure to use the random number generator defined in the setup method\n",
    "        n_electron_interface = n_electron_interface * self.rng.normal(\n",
    "            1, 0.1, len(n_electron_interface)\n",
    "        )\n",
    "        result[\"n_electron_interface\"][mask] = n_electron_interface\n",
    "\n",
    "        return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that our new plugin is defined we can register it, adjust the config and then try to run it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st = fuse.context.xenonnt_fuse_full_chain_simulation(\n",
    "    output_folder=\"./fuse_data\",\n",
    "    simulation_config=\"sr0_dev\",\n",
    ")\n",
    "\n",
    "# We will use a different file now with a little more statistics\n",
    "st.set_config(\n",
    "    {\n",
    "        \"path\": (\n",
    "            \"/project2/lgrandi/xenonnt/simulations/lead-214/high_energy_er_fullchain_Pb214_simulations/TPC_Pb214_lower/G4outsTPCXenonProgeny1_2000_nochain_GS_42/\"\n",
    "        ),\n",
    "        \"file_name\": \"nT_TPC_Pb214_2000_nochain_50.root\",\n",
    "        \"n_periods\": 5,\n",
    "    }\n",
    ")\n",
    "\n",
    "run_number = \"00000\"\n",
    "\n",
    "st.register(CosineElectronDrift)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st.make(run_number, \"drifted_electrons\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = st.get_df(run_number, [\"microphysics_summary\", \"drifted_electrons\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that our new plugin produced some data we can try to visualize the simulation output. We will calculate the mean n_electron_interface for different z-slices and plot the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def average_value_in_interval(edge_start, edge_end, df, value):\n",
    "    value_mean_in_interval = []\n",
    "    for begin, end in zip(edge_start, edge_end):\n",
    "        data_in_slice = df[(df.z >= begin) & (df.z < end)][value]\n",
    "\n",
    "        if len(data_in_slice) > 0:\n",
    "            value_mean_in_interval.append(np.nanmean(data_in_slice))\n",
    "\n",
    "        else:\n",
    "            value_mean_in_interval.append(0)\n",
    "\n",
    "    value_mean_in_interval = np.array(value_mean_in_interval)\n",
    "\n",
    "    return value_mean_in_interval\n",
    "\n",
    "\n",
    "bin_edges = np.linspace(-150, -1, 150)\n",
    "bin_edges_start = bin_edges[:-1]\n",
    "bin_edges_end = bin_edges[1:]\n",
    "bin_centers = (bin_edges_start + bin_edges_end) / 2\n",
    "\n",
    "electrons_reaching_the_interface_in_z_slice = average_value_in_interval(\n",
    "    bin_edges_start, bin_edges_end, data, \"n_electron_interface\"\n",
    ")\n",
    "electrons_at_interaction_site_in_z_slice = average_value_in_interval(\n",
    "    bin_edges_start, bin_edges_end, data, \"electrons\"\n",
    ")\n",
    "\n",
    "\n",
    "plt.plot(\n",
    "    bin_centers,\n",
    "    electrons_reaching_the_interface_in_z_slice,\n",
    "    color=\"purple\",\n",
    "    label=\"Electrons reaching the interface\",\n",
    "    lw=0,\n",
    "    marker=\"o\",\n",
    "    markersize=3,\n",
    ")\n",
    "\n",
    "plt.plot(\n",
    "    bin_centers,\n",
    "    electrons_at_interaction_site_in_z_slice,\n",
    "    color=\"orange\",\n",
    "    label=\"Electrons at interaction site\",\n",
    "    lw=0,\n",
    "    marker=\"o\",\n",
    "    markersize=3,\n",
    ")\n",
    "\n",
    "plt.legend()\n",
    "plt.xlabel(\"z [cm]\")\n",
    "plt.ylabel(\"Number of electrons\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tipps & Tricks for plugin development\n",
    "\n",
    "In this section I will collect some tipps and tricks that might be useful when developing new plugins. \n",
    "\n",
    "#### 1. Running the compute method in the notebook\n",
    "\n",
    "strax has the functionality to access a plugin directly. This way you can test your plugins methods in the notebook without relying on `st.make` to run the plugin. First we can get the plugin completely initialized by calling `st.get_single_plugin`. Then we can call the `compute` method with the needed input arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plugin = st.get_single_plugin(run_number, \"drifted_electrons\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make sure that your input data has the correct numpy format. Dataframes are not supported."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "microphysics_summary = st.get_array(run_number, [\"microphysics_summary\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plugin_output = plugin.compute(microphysics_summary)\n",
    "print(plugin_output[0:10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Replacing a map with a dummy map\n",
    "\n",
    "Sometimes you might want to disable a certain feature of the simulation without going into the code and changing a plugin. One example can be that you might want to turn of the effects of a certain map. In this example we will use the `constant_dummy_map` URLConfig protocol to disable the effect of the s2_correction map.\n",
    "\n",
    "We can set up a new context. This time we will override the config option for the s2_correction map with the dummy map option. Instead of the real map values, the dummy map will return a constant value that can be configured by the user. In our case we set it to 1. You could also set it to other values if you want to test the effect of the map on the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st = fuse.context.xenonnt_fuse_full_chain_simulation(\n",
    "    output_folder=\"./fuse_data\",\n",
    "    simulation_config=\"sr0_dev\",\n",
    ")\n",
    "\n",
    "st.set_config(\n",
    "    {\n",
    "        \"path\": \"/project2/lgrandi/xenonnt/simulations/testing\",\n",
    "        \"file_name\": \"pmt_neutrons_100.root\",\n",
    "        \"entry_stop\": 10,\n",
    "        \"s2_correction_map\": \"constant_dummy_map://1\",\n",
    "    }\n",
    ")\n",
    "\n",
    "run_number = \"00000\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get a feeling how the map affects your simulation, lets take a look at the number of s2 photons.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st.make(run_number, \"s2_photons\")\n",
    "s2_photons = st.get_df(run_number, \"s2_photons\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the data, lets change the map to a different value and rerun the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st.set_config(\n",
    "    {\n",
    "        \"s2_correction_map\": \"constant_dummy_map://2\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "st.make(run_number, \"s2_photons\")\n",
    "s2_photons_changed_constant = st.get_df(run_number, \"s2_photons\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we can plot the data again. You should see that the distribution of s2 photons is different."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "bins = np.linspace(0, 100, 100)\n",
    "\n",
    "plt.hist(s2_photons[\"n_s2_photons\"], bins=bins, alpha=0.5, label=\"Original\")\n",
    "plt.hist(\n",
    "    s2_photons_changed_constant[\"n_s2_photons\"], bins=bins, alpha=0.5, label=\"Changed constant\"\n",
    ")\n",
    "\n",
    "plt.legend()\n",
    "plt.xlabel(\"Number of S2 photons\")\n",
    "plt.ylabel(\"Counts\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.19"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}