Add option to change source distance

src/tof/source.py

@@ -454,6 +454,8 @@ class Source:
         List of choppers to optimize the source for. A chopper acceptance diagram will
         be overlaid on the source distribution, and samples will be taken only from
         the regions where the chopper cascade is accepting neutrons.
+    distance:
+        Position of the source along the beamline.
 
         .. versionadded:: 26.4.0
     """  # noqa: E501
@@ -470,6 +472,7 @@ def __init__(
         tmax: sc.Variable | None = None,
         seed: int | None = None,
         optimize_for: list[Chopper] | None = None,
+        distance: sc.Variable | None = None,
     ):
         self._facility = facility.lower() if facility is not None else None
         self._neutrons = int(neutrons)
@@ -485,7 +488,9 @@ def __init__(
         if frequency is not None:
             facility_pulse = facility_pulse.assign_coords(frequency=frequency)
 
-        self._distance = facility_pulse.coords['distance']
+        self._distance = (
+            facility_pulse.coords['distance'] if distance is None else distance
+        )
         self._frequency = facility_pulse.coords['frequency']
         self._data = _sample_source_data_from_distribution(
             p=facility_pulse,

tests/source_test.py

@@ -47,6 +47,12 @@ def test_ess_pulse_wmin_wmax():
     assert sc.all(source.data['pulse', 0].coords['wavelength'] <= wmax)
 
 
+def test_ess_pulse_with_distance():
+    distance = sc.scalar(2.35, unit='m')
+    source = tof.Source(facility='ess-odin', neutrons=100_000, distance=distance)
+    assert sc.identical(source.distance, distance)
+
+
 def test_creation_from_neutrons():
     birth_times = sc.array(dims=['event'], values=[1000.0, 1500.0, 2000.0], unit='us')
     wavelengths = sc.array(dims=['event'], values=[1.0, 5.0, 10.0], unit='angstrom')
@@ -59,6 +65,18 @@ def test_creation_from_neutrons():
     assert sc.identical(source.data['pulse', 0].coords['wavelength'], wavelengths)
 
 
+def test_from_neutrons_with_distance():
+    distance = sc.scalar(2.35, unit='m')
+    birth_times = sc.array(dims=['event'], values=[1000.0, 1500.0, 2000.0], unit='us')
+    wavelengths = sc.array(dims=['event'], values=[1.0, 5.0, 10.0], unit='angstrom')
+    source = tof.Source.from_neutrons(
+        birth_times=birth_times,
+        wavelengths=wavelengths,
+        distance=distance,
+    )
+    assert sc.identical(source.distance, distance)
+
+
 def test_creation_from_distribution_flat():
     birth_time = sc.linspace('birth_time', 1.0, 3.0, 100, unit='ms')
     p_time = sc.DataArray(
@@ -88,6 +106,30 @@ def test_creation_from_distribution_flat():
     )
 
 
+def test_from_distribution_with_distance():
+    distance = sc.scalar(2.35, unit='m')
+    birth_time = sc.linspace('birth_time', 1.0, 3.0, 100, unit='ms')
+    p_time = sc.DataArray(
+        data=sc.ones(sizes=birth_time.sizes),
+        coords={'birth_time': birth_time},
+    )
+    wavelength = sc.linspace('wavelength', 1.0, 10.0, 100, unit='angstrom')
+    p_wav = sc.DataArray(
+        data=sc.ones(sizes=wavelength.sizes),
+        coords={'wavelength': wavelength},
+    )
+
+    N = 123456
+    source = tof.Source.from_distribution(
+        neutrons=N,
+        p_time=p_time,
+        p_wav=p_wav,
+        distance=distance,
+    )
+
+    assert sc.identical(source.distance, distance)
+
+
 @pytest.mark.parametrize('distribution_2d', [False, True])
 def test_creation_from_distribution(distribution_2d):
     v = np.ones(90) * 0.1