Time binning reduction and export method.

docs/api-reference/index.md

@@ -11,7 +11,9 @@
    :recursive:
 
    NMXData
+   NMXReducedData
    InputFilepath
+   TimeBinSteps
 
 ```
 

docs/examples/workflow.ipynb

@@ -26,20 +26,42 @@
     "from ess.nmx.mcstas_loader import (\n",
     "    InputFilepath,\n",
     "    MaximumProbability,\n",
-    "    DefaultMaximumProbability\n",
+    "    DefaultMaximumProbability,\n",
+    "    EventWeightsConverter,\n",
+    "    event_weights_from_probability,\n",
+    "    ProtonChargeConverter,\n",
+    "    proton_charge_from_event_data,\n",
     ")\n",
     "from ess.nmx.data import small_mcstas_sample\n",
+    "from ess.nmx.reduction import bin_time_of_arrival, TimeBinSteps\n",
     "\n",
-    "providers = (load_mcstas_nexus, )\n",
+    "providers = (load_mcstas_nexus, bin_time_of_arrival, )\n",
     "\n",
     "file_path = small_mcstas_sample()  # Replace it with your data file path\n",
     "params = {\n",
+    "    TimeBinSteps: TimeBinSteps(50),\n",
     "    InputFilepath: InputFilepath(file_path),\n",
     "    # Additional parameters for McStas data handling.\n",
     "    MaximumProbability: DefaultMaximumProbability,\n",
+    "    EventWeightsConverter: event_weights_from_probability,\n",
+    "    ProtonChargeConverter: proton_charge_from_event_data,\n",
     "}"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "``event weights converter`` and ``proton_charge_from_event_data`` are\n",
+    "\n",
+    "set as parameters for reproducibility of workflow and accessibility to the documentation.\n",
+    "\n",
+    "The reason of having them as parameters not as providers is,\n",
+    "\n",
+    "1. They are not part of general reduction, which are only for McStas cases.\n",
+    "2. They are better done while the file is open and read in the loader.\n"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -71,9 +93,10 @@
    "source": [
     "import sciline as sl\n",
     "from ess.nmx.mcstas_loader import NMXData\n",
+    "from ess.nmx.reduction import NMXReducedData\n",
     "\n",
     "nmx_pl = sl.Pipeline(list(providers), params=params)\n",
-    "nmx_workflow = nmx_pl.get(NMXData)\n",
+    "nmx_workflow = nmx_pl.get(NMXReducedData)\n",
     "nmx_workflow.visualize()"
    ]
   },
@@ -91,8 +114,35 @@
    "outputs": [],
    "source": [
     "# Event data grouped by detector panel and pixel id.\n",
-    "da = nmx_workflow.compute(NMXData)\n",
-    "da"
+    "dg = nmx_workflow.compute(NMXData)\n",
+    "dg"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Binned data.\n",
+    "\n",
+    "binned_dg = nmx_workflow.compute(NMXReducedData)\n",
+    "binned_dg"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Export Results\n",
+    "\n",
+    "``NMXReducedData`` object has a method to export the data into nexus or h5 file.\n",
+    "\n",
+    "You can save the result as ``test.nxs`` for example.\n",
+    "\n",
+    "```python\n",
+    "binned_dg.export_as_nexus('test.nxs')\n",
+    "```"
    ]
   },
   {
@@ -107,7 +157,7 @@
     "All pixel positions are relative to the sample position,\n",
     "therefore the sample is at (0, 0, 0).\n",
     "\n",
-    "The instrument view is shown using"
+    "**It might be very slow or not work in the ``VS Code`` jupyter notebook editor.**"
    ]
   },
   {
@@ -118,8 +168,12 @@
    "source": [
     "import scippneutron as scn\n",
     "\n",
+    "da = dg['weights']\n",
+    "da.coords['position'] = dg['position']\n",
     "# Plot one out of 100 pixels to reduce size of docs output\n",
-    "scn.instrument_view(da['id', ::100].hist(), pixel_size=0.0075)"
+    "view = scn.instrument_view(da['id', ::100].hist(), pixel_size=0.0075)\n",
+    "view.children[0].toolbar.cameraz()\n",
+    "view"
    ]
   }
  ],
@@ -138,7 +192,8 @@
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3"
+   "pygments_lexer": "ipython3",
+   "version": "3.9.18"
   }
  },
  "nbformat": 4,

src/ess/nmx/__init__.py

@@ -13,10 +13,13 @@
 
 from .data import small_mcstas_sample
 from .mcstas_loader import InputFilepath, NMXData, load_mcstas_nexus
+from .reduction import NMXReducedData, TimeBinSteps
 
 __all__ = [
     "small_mcstas_sample",
     "NMXData",
     "InputFilepath",
     "load_mcstas_nexus",
+    "NMXReducedData",
+    "TimeBinSteps",
 ]

src/ess/nmx/mcstas_loader.py

@@ -1,18 +1,33 @@
 # SPDX-License-Identifier: BSD-3-Clause
 # Copyright (c) 2023 Scipp contributors (https://github.com/scipp)
-from typing import Iterable, NewType, Optional
+from typing import Callable, Iterable, NewType, Optional
 
 import scipp as sc
 import scippnexus as snx
 
+from .reduction import NMXData
+
 PixelIDs = NewType("PixelIDs", sc.Variable)
 InputFilepath = NewType("InputFilepath", str)
-NMXData = NewType("NMXData", sc.DataGroup)
 
 # McStas Configurations
 MaximumProbability = NewType("MaximumProbability", int)
 DefaultMaximumProbability = MaximumProbability(100_000)
 
+McStasEventProbabilities = NewType("McStasEventProbabilities", sc.Variable)
+EventWeights = NewType("EventWeights", sc.Variable)
+EventWeightsConverter = NewType(
+    "EventWeightsConverter",
+    Callable[[MaximumProbability, McStasEventProbabilities], EventWeights],
+)
+"""A function that converts McStas probability to event weights."""
+
+ProtonCharge = NewType("ProtonCharge", sc.Variable)
+ProtonChargeConverter = NewType(
+    "ProtonChargeConverter", Callable[[EventWeights], ProtonCharge]
+)
+"""A function that derives arbitrary proton charge based on event weights."""
+
 
 def _retrieve_event_list_name(keys: Iterable[str]) -> str:
     prefix = "bank01_events_dat_list"
@@ -27,54 +42,172 @@ def _retrieve_event_list_name(keys: Iterable[str]) -> str:
     raise ValueError("Can not find event list name.")
 
 
+def _retrieve_raw_event_data(file: snx.File) -> sc.Variable:
+    """Retrieve events from the nexus file."""
+    bank_name = _retrieve_event_list_name(file["entry1/data"].keys())
+    # ``dim_0``: event index, ``dim_1``: property index.
+    return file["entry1/data/" + bank_name]["events"][()].rename_dims(
+        {'dim_0': 'event'}
+    )
+
+
 def _copy_partial_var(
     var: sc.Variable, idx: int, unit: Optional[str] = None, dtype: Optional[str] = None
 ) -> sc.Variable:
     """Retrieve a property from a variable."""
     var = var['dim_1', idx].astype(dtype or var.dtype, copy=True)
-    if unit:
+    if unit is not None:
         var.unit = sc.Unit(unit)
     return var
 
 
+def _retrieve_crystal_rotation(file: snx.File, unit: str) -> sc.Variable:
+    """Retrieve crystal rotation from the file."""
+
+    return sc.vector(
+        value=[file[f"entry1/simulation/Param/XtalPhi{key}"][...] for key in "XYZ"],
+        unit=unit,
+    )
+
+
+def event_weights_from_probability(
+    max_probability: MaximumProbability, probabilities: McStasEventProbabilities
+) -> EventWeights:
+    """Create event weights by scaling probability data.
+
+    event_weights = max_probability * (probabilities / max(probabilities))
+
+    Parameters
+    ----------
+    probabilities:
+        The probabilities of the events.
+
+    max_probability:
+        The maximum probability to scale the weights.
+
+    """
+    maximum_probability = sc.scalar(max_probability, unit='counts')
+
+    return EventWeights(maximum_probability * (probabilities / probabilities.max()))
+
+
+def _compose_event_data_array(
+    *,
+    weights: sc.Variable,
+    id_list: sc.Variable,
+    t_list: sc.Variable,
+    pixel_ids: sc.Variable,
+    num_panels: int,
+) -> sc.DataArray:
+    """Combine data with coordinates loaded from the nexus file.
+
+    Parameters
+    ----------
+    weights:
+        The weights of the events.
+
+    id_list:
+        The pixel IDs of the events.
+
+    t_list:
+        The time of arrival of the events.
+
+    pixel_ids:
+        All possible pixel IDs of the detector.
+
+    num_panels:
+        The number of (detector) panels used in the experiment.
+
+    """
+
+    events = sc.DataArray(data=weights, coords={'t': t_list, 'id': id_list})
+    grouped: sc.DataArray = events.group(pixel_ids)
+    return grouped.fold(dim='id', sizes={'panel': num_panels, 'id': -1})
+
+
+def proton_charge_from_event_data(event_da: sc.DataArray) -> ProtonCharge:
+    """Make up the proton charge from the event data array.
+
+    Proton charge is proportional to the number of neutrons,
+    which is proportional to the number of events.
+    The scale factor is manually chosen based on previous results
+    to be convenient for data manipulation in the next steps.
+    It is derived this way since
+    the protons are not part of McStas simulation,
+    and the number of neutrons is not included in the result.
+
+    Parameters
+    ----------
+    event_da:
+        The event data binned in detector panel and pixel id dimensions.
+
+    """
+    # Arbitrary number to scale the proton charge
+    _proton_charge_scale_factor = sc.scalar(1 / 10_000, unit=None)
+
+    return ProtonCharge(_proton_charge_scale_factor * event_da.bins.size().sum().data)
+
+
 def load_mcstas_nexus(
     file_path: InputFilepath,
+    event_weights_converter: EventWeightsConverter = event_weights_from_probability,
+    proton_charge_converter: ProtonChargeConverter = proton_charge_from_event_data,
     max_probability: Optional[MaximumProbability] = None,
 ) -> NMXData:
     """Load McStas simulation result from h5(nexus) file.
 
+    See :func:`~event_weights_from_probability` and
+    :func:`~proton_charge_from_event_data` for details.
+
     Parameters
     ----------
     file_path:
         File name to load.
 
+    event_weights_converter:
+        A function to convert probabilities to event weights.
+        The function should accept the probabilities as the first argument,
+        and return the converted event weights.
+
+    proton_charge_converter:
+        A function to convert the event weights to proton charge.
+        The function should accept the event weights as the first argument,
+        and return the proton charge.
+
     max_probability:
         The maximum probability to scale the weights.
+        If not provided, ``DefaultMaximumProbability`` is used.
 
     """
 
     from .mcstas_xml import read_mcstas_geometry_xml
 
     geometry = read_mcstas_geometry_xml(file_path)
-    probability = max_probability or DefaultMaximumProbability
+    coords = geometry.to_coords()
 
     with snx.File(file_path) as file:
-        bank_name = _retrieve_event_list_name(file["entry1/data"].keys())
-        var: sc.Variable
-        var = file["entry1/data/" + bank_name]["events"][()].rename_dims(
-            {'dim_0': 'event'}
+        raw_data = _retrieve_raw_event_data(file)
+        weights = event_weights_converter(
+            max_probability or DefaultMaximumProbability,
+            McStasEventProbabilities(
+                _copy_partial_var(raw_data, idx=0, unit='counts')
+            ),  # p
+        )
+        event_da = _compose_event_data_array(
+            weights=weights,
+            id_list=_copy_partial_var(raw_data, idx=4, dtype='int64'),  # id
+            t_list=_copy_partial_var(raw_data, idx=5, unit='s'),  # t
+            pixel_ids=coords.pop('pixel_id'),
+            num_panels=len(geometry.detectors),
+        )
+        proton_charge = proton_charge_converter(event_da)
+        crystal_rotation = _retrieve_crystal_rotation(
+            file, geometry.simulation_settings.angle_unit
         )
 
-        weights = _copy_partial_var(var, idx=0, unit='counts')  # p
-        id_list = _copy_partial_var(var, idx=4, dtype='int64')  # id
-        t_list = _copy_partial_var(var, idx=5, unit='s')  # t
-
-        weights = (probability / weights.max()) * weights
-
-        loaded = sc.DataArray(data=weights, coords={'t': t_list, 'id': id_list})
-        coords = geometry.to_coords()
-        grouped = loaded.group(coords.pop('pixel_id'))
-        da = grouped.fold(dim='id', sizes={'panel': len(geometry.detectors), 'id': -1})
-        da.coords.update(coords)
-
-        return NMXData(da)
+    return NMXData(
+        weights=event_da,
+        proton_charge=proton_charge,
+        crystal_rotation=crystal_rotation,
+        **coords,
+    )