essreduce 25.1.0__py3-none-any.whl → 25.2.0__py3-none-any.whl

ess/reduce/time_of_flight/fakes.py

@@ -21,6 +21,7 @@ class FakeBeamline:
         monitors: dict[str, sc.Variable],
         run_length: sc.Variable,
         events_per_pulse: int = 200000,
+        seed: int | None = None,
         source: Callable | None = None,
     ):
         import math
@@ -35,7 +36,10 @@ class FakeBeamline:
         # Create a source
         if source is None:
             self.source = tof_pkg.Source(
-                facility="ess", neutrons=self.events_per_pulse, pulses=self.npulses
+                facility="ess",
+                neutrons=self.events_per_pulse,
+                pulses=self.npulses,
+                seed=seed,
             )
         else:
             self.source = source(pulses=self.npulses)
@@ -69,34 +73,10 @@ class FakeBeamline:
         self.model_result = self.model.run()
 
     def get_monitor(self, name: str) -> sc.DataGroup:
-        # Create some fake pulse time zero
-        start = sc.datetime("2024-01-01T12:00:00.000000")
-        period = sc.reciprocal(self.frequency)
-
-        detector = self.model_result.detectors[name]
-        raw_data = detector.data.flatten(to="event")
-        # Select only the neutrons that make it to the detector
+        nx_event_data = self.model_result.to_nxevent_data(name)
+        raw_data = self.model_result.detectors[name].data.flatten(to="event")
         raw_data = raw_data[~raw_data.masks["blocked_by_others"]].copy()
-        raw_data.coords["Ltotal"] = detector.distance
-
-        # Format the data in a way that resembles data loaded from NeXus
-        event_data = raw_data.copy(deep=False)
-        dt = period.to(unit="us")
-        event_time_zero = (dt * (event_data.coords["toa"] // dt)).to(dtype=int) + start
-        raw_data.coords["event_time_zero"] = event_time_zero
-        event_data.coords["event_time_zero"] = event_time_zero
-        event_data.coords["event_time_offset"] = (
-            event_data.coords.pop("toa").to(unit="s") % period
-        )
-        del event_data.coords["tof"]
-        del event_data.coords["speed"]
-        del event_data.coords["time"]
-        del event_data.coords["wavelength"]
-
-        return (
-            event_data.group("event_time_zero").rename_dims(event_time_zero="pulse"),
-            raw_data.group("event_time_zero").rename_dims(event_time_zero="pulse"),
-        )
+        return nx_event_data, raw_data
 
 
 wfm1_chopper = DiskChopper(
@@ -194,25 +174,6 @@ pol_chopper = DiskChopper(
     radius=sc.scalar(30.0, unit="cm"),
 )
 
-pulse_skipping = DiskChopper(
-    frequency=sc.scalar(-7.0, unit="Hz"),
-    beam_position=sc.scalar(0.0, unit="deg"),
-    phase=sc.scalar(0.0, unit="deg"),
-    axle_position=sc.vector(value=[0, 0, 30.0], unit="m"),
-    slit_begin=sc.array(
-        dims=["cutout"],
-        values=np.array([40.0]),
-        unit="deg",
-    ),
-    slit_end=sc.array(
-        dims=["cutout"],
-        values=np.array([140.0]),
-        unit="deg",
-    ),
-    slit_height=sc.scalar(10.0, unit="cm"),
-    radius=sc.scalar(30.0, unit="cm"),
-)
-
 
 def wfm_choppers():
     return {
@@ -238,3 +199,24 @@ def psc_choppers():
         )
         for name, ch in wfm_choppers().items()
     }
+
+
+def pulse_skipping_chopper():
+    return DiskChopper(
+        frequency=sc.scalar(-7.0, unit="Hz"),
+        beam_position=sc.scalar(0.0, unit="deg"),
+        phase=sc.scalar(0.0, unit="deg"),
+        axle_position=sc.vector(value=[0, 0, 30.0], unit="m"),
+        slit_begin=sc.array(
+            dims=["cutout"],
+            values=np.array([40.0]),
+            unit="deg",
+        ),
+        slit_end=sc.array(
+            dims=["cutout"],
+            values=np.array([140.0]),
+            unit="deg",
+        ),
+        slit_height=sc.scalar(10.0, unit="cm"),
+        radius=sc.scalar(30.0, unit="cm"),
+    )

ess/reduce/time_of_flight/simulation.py

@@ -11,6 +11,7 @@ from .types import SimulationResults
 def simulate_beamline(
     choppers: Mapping[str, DiskChopper],
     neutrons: int = 1_000_000,
+    pulses: int = 1,
     seed: int | None = None,
     facility: str = 'ess',
 ) -> SimulationResults:
@@ -26,6 +27,8 @@ def simulate_beamline(
         for more information.
     neutrons:
         Number of neutrons to simulate.
+    pulses:
+        Number of pulses to simulate.
     seed:
         Seed for the random number generator used in the simulation.
     facility:
@@ -47,9 +50,9 @@ def simulate_beamline(
         )
         for name, ch in choppers.items()
     ]
-    source = tof.Source(facility=facility, neutrons=neutrons, seed=seed)
+    source = tof.Source(facility=facility, neutrons=neutrons, pulses=pulses, seed=seed)
     if not tof_choppers:
-        events = source.data.squeeze()
+        events = source.data.squeeze().flatten(to='event')
         return SimulationResults(
             time_of_arrival=events.coords["time"],
             speed=events.coords["speed"],
@@ -61,7 +64,7 @@ def simulate_beamline(
     results = model.run()
     # Find name of the furthest chopper in tof_choppers
     furthest_chopper = max(tof_choppers, key=lambda c: c.distance)
-    events = results[furthest_chopper.name].data.squeeze()
+    events = results[furthest_chopper.name].data.squeeze().flatten(to='event')
     events = events[
         ~(events.masks["blocked_by_others"] | events.masks["blocked_by_me"])
     ]

ess/reduce/time_of_flight/to_events.py

@@ -34,30 +34,37 @@ def to_events(
     rng = np.random.default_rng()
     event_coords = {}
     edge_dims = []
-    midp_dims = []
+    midp_dims = set()
+    midp_coord_names = []
     # Separate bin-edge and midpoints coords
-    for dim in da.dims:
-        if da.coords.is_edges(dim):
-            edge_dims.append(dim)
+    for name in da.coords:
+        dims = da.coords[name].dims
+        is_edges = False if not dims else da.coords.is_edges(name)
+        if is_edges:
+            if name in dims:
+                edge_dims.append(name)
         else:
-            midp_dims.append(dim)
+            midp_coord_names.append(name)
+            midp_dims.update(set(dims))
 
-    edge_sizes = {dim: da.sizes[dim] for dim in edge_dims}
+    edge_sizes = {dim: da.sizes[da.coords[dim].dim] for dim in edge_dims}
     for dim in edge_dims:
         coord = da.coords[dim]
-        low = sc.broadcast(coord[dim, :-1], sizes=edge_sizes).values
-        high = sc.broadcast(coord[dim, 1:], sizes=edge_sizes).values
+        left = sc.broadcast(coord[dim, :-1], sizes=edge_sizes).values
+        right = sc.broadcast(coord[dim, 1:], sizes=edge_sizes).values
 
         # The numpy.random.uniform function below does not support NaNs, so we need to
         # replace them with zeros, and then replace them back after the random numbers
         # have been generated.
-        nans = np.isnan(low) | np.isnan(high)
-        low = np.where(nans, 0.0, low)
-        high = np.where(nans, 0.0, high)
+        nans = np.isnan(left) | np.isnan(right)
+        left = np.where(nans, 0.0, left)
+        right = np.where(nans, 0.0, right)
+        # Ensure left <= right
+        left, right = np.minimum(left, right), np.maximum(left, right)
 
         # In each bin, we generate a number of events with a uniform distribution.
         events = rng.uniform(
-            low, high, size=(events_per_bin, *list(edge_sizes.values()))
+            left, right, size=(events_per_bin, *list(edge_sizes.values()))
         )
         events[..., nans] = np.nan
         event_coords[dim] = sc.array(
@@ -77,20 +84,20 @@ def to_events(
     data = da.data
     if event_masks:
         inv_mask = (~reduce(lambda a, b: a | b, event_masks.values())).to(dtype=int)
-        inv_mask.unit = ''
+        inv_mask.unit = ""
         data = data * inv_mask
 
     # Create the data counts, which are the original counts divided by the number of
     # events per bin
     sizes = {event_dim: events_per_bin} | da.sizes
     val = sc.broadcast(sc.values(data) / float(events_per_bin), sizes=sizes)
-    kwargs = {'dims': sizes.keys(), 'values': val.values, 'unit': data.unit}
+    kwargs = {"dims": sizes.keys(), "values": val.values, "unit": data.unit}
     if data.variances is not None:
         # Note here that all the events are correlated.
         # If we later histogram the events with different edges than the original
         # histogram, then neighboring bins will be correlated, and the error obtained
         # will be too small. It is however not clear what can be done to improve this.
-        kwargs['variances'] = sc.broadcast(
+        kwargs["variances"] = sc.broadcast(
             sc.variances(data) / float(events_per_bin), sizes=sizes
         ).values
     new_data = sc.array(**kwargs)
@@ -100,5 +107,5 @@ def to_events(
         dims=[*edge_dims, event_dim], to=event_dim
     )
     return new.assign_coords(
-        {dim: da.coords[dim].copy() for dim in midp_dims}
+        {dim: da.coords[dim].copy() for dim in midp_coord_names}
     ).assign_masks({key: mask.copy() for key, mask in other_masks.items()})