[BEER] Time-of-flight for pulse shaping mode

[jokasimr]

The current Beer McStas file loader can only handle multiplexing chopper modes.

It should also be able to handle pulse shaping modes.

Celine already has a script that does this.

[celinedurniak]

FYI Modified version of io.py to consider the pulse shaping chopper modes (to be tested for all modes)

# io.py

from ess.beer.types import (
    DetectorData,
    Filename,
    ModulationPeriod,
    SampleRun,
    TwoThetaMaskFunction,
    WavelengthDefinitionChopperDelay,
)

def _load_h5(group: h5py.Group | str, *paths: str):
    if isinstance(group, str):
        with h5py.File(group) as group:
            yield from _load_h5(group, *paths)
        return
    for path in paths:
        if type(path) is str:
            g = group
            
            for p in path.strip('/').split('/'):
                g = _unique_child_group_h5(g, p) if p.startswith('NX') else g.get(p)
        elif type(path) is list:
            pass
        yield g


def _unique_child_group_h5(
    group: h5py.Group,
    nx_class: str,
) -> h5py.Group | None:
    out = None
    for v in group.values():
        if v.attrs.get("NX_class") == nx_class.encode():
            if out is None:
                out = v
            else:
                raise ValueError(
                    f'Expected exactly one {nx_class} group, but found more'
                )
    return out


def _load_beer_mcstas(f, bank=1):
    for key in f['/entry1/instrument/components']:
        if 'sampleMantid' in key:
            sample_position_path = f'/entry1/instrument/components/{key}/Position'
            break
        elif 'PSC1' in key:
            psc1_position_path = f'/entry1/instrument/components/{key}/Position'
        elif 'PSC2' in key:
            psc2_position_path = f'/entry1/instrument/components/{key}/Position'
        elif 'PSC3' in key:
            psc3_position_path = f'/entry1/instrument/components/{key}/Position'
        elif 'MCA' in key:
            mca_position_path = f'/entry1/instrument/components/{key}/Position'
        elif 'MCC' in key:
            mcc_position_path = f'/entry1/instrument/components/{key}/Position'
    else:
        raise ValueError('Sample or choppers position entries not found in file.')

    #data, events, params, sample_pos, chopper_pos = _load_h5(
    data, events, params, sample_pos, psc1_pos, psc2_pos, psc3_pos, mca_pos, mcc_pos = _load_h5(
        f,
        f'NXentry/NXdetector/bank{bank:02}_events_dat_list_p_x_y_n_id_t',
        f'NXentry/NXdetector/bank{bank:02}_events_dat_list_p_x_y_n_id_t/events',
        'NXentry/simulation/Param',
        sample_position_path,
        psc1_position_path,
        psc2_position_path,
        psc3_position_path,
        mca_position_path,
        mcc_position_path,
    )
    
    events = events[()]
    da = sc.DataArray(
        sc.array(dims=['events'], values=events[:, 0], variances=events[:, 0] ** 2),
    )
    for name, value in data.attrs.items():
        if name in ('position',):
            da.coords[name] = sc.scalar(value.decode())

    for i, label in enumerate(data.attrs["ylabel"].decode().strip().split(' ')):
        if label == 'p':
            continue
        da.coords[label] = sc.array(dims=['events'], values=events[:, i])

    for k, v in params.items():
        v = v[0]
        if isinstance(v, bytes):
            v = v.decode()
        if k in ('mode', 'sample_filename', 'lambda'):
            da.coords[k] = sc.scalar(v)

    if da.coords['lambda'].value == '0':
        if da.coords['mode'].value in ['0', '3', '4', '5', '6', '7', '8', '9', '10', '15', '16']:
            da.coords['lambda'].value =  '2.1'
        elif da.coords['mode'].value in ['1', '2']:
            da.coords['lambda'].value = '3.1'
        elif da.coords['mode'].value == '11':
            da.coords['lambda'].value = '3'
        elif da.coords['mode'].value == '12':	
            da.coords['lambda'].value = '3.5'
        elif da.coords['mode'].value == '13':
            da.coords['lambda'].value = '6'
        elif da.coords['mode'].value == '14':
            da.coords['lambda'].value = '4'

    da.coords['sample_position'] = sc.vector(sample_pos[:], unit='m')
    da.coords['detector_position'] = sc.vector(
        list(map(float, da.coords.pop('position').value.split(' '))), unit='m'
    )
        
    if da.coords['mode'].value in ['0', '1', '2', '11', '16']:
        chopper_pos = np.array([0., 0., 0.])
    elif da.coords['mode'].value in ['3', '4', '12', '13', '15']:
        chopper_pos = 0.5 * (psc1_pos[:] + psc3_pos[:])
    elif da.coords['mode'].value in ['5', '6']:
        chopper_pos = 0.5 * (psc1_pos[:] + psc2_pos[:])
    elif da.coords['mode'].value in ['7', '8', '9', '10']:
        chopper_pos = mca_pos[:]
    elif da.coords['mode'].value  == '14':
        chopper_pos = mcc_pos[:]

    da.coords['chopper_position'] = sc.vector(chopper_pos[:], unit='m')
    da.coords['x'].unit = 'm'
    da.coords['y'].unit = 'm'
    da.coords['t'].unit = 's'

    z = sc.norm(da.coords['detector_position'] - da.coords['sample_position'])
    # z = da.coords['position'].fields.z - da.coords['sample_position'].fields.z
    L1 = sc.norm(da.coords['sample_position'] - da.coords['chopper_position'])
    L2 = sc.sqrt(da.coords['x'] ** 2 + da.coords['y'] ** 2 + z**2)
    # Source is assumed to be at the origin
    
    da.coords['L0'] = L1 + L2 + sc.norm(da.coords['chopper_position'])
    da.coords['Ltotal'] = L1 + L2
    da.coords['two_theta'] = sc.acos(
        (-da.coords['x'] if bank == 1 else da.coords['x']) / L2
    )

    # Save some space
    da.coords.pop('x')
    da.coords.pop('y')
    da.coords.pop('n')

    # expression of temporal offset delta_t checked for pulse shaping mode: 4,5,6
    delta_t =  m_n.value / h.value * float(da.coords['lambda'].value) * sc.norm(da.coords['chopper_position']).value * 1e-10 * sc.Unit('s')
    # print(da.coords['mode'].value, delta_t)
    
    da.coords["t"] -= delta_t
    # da.coords['event_time_offset'] = da.coords.pop('t')
    da.coords['tof'] = da.coords.pop('t')
    return da


def load_beer_mcstas(f: str | Path | h5py.File) -> sc.DataGroup:
    '''Load beer McStas data from a file to a
    data group with one data array for each bank.
    '''
    if isinstance(f, str | Path):
        with h5py.File(f) as ff:
            return load_beer_mcstas(ff)

    return sc.DataGroup(
        {
            'bank1': _load_beer_mcstas(f, bank=1),
            'bank2': _load_beer_mcstas(f, bank=2),
        }
    )


def load_beer_mcstas_provider(
    fname: Filename[SampleRun], two_theta_mask: TwoThetaMaskFunction
) -> DetectorData[SampleRun]:
    da = load_beer_mcstas(fname)
    da = (
        sc.DataGroup(
            {
                k: v.assign_masks(two_theta=two_theta_mask(v.coords['two_theta']))
                for k, v in da.items()
            }
        )
        if isinstance(da, sc.DataGroup)
        else da.assign_masks(two_theta=two_theta_mask(da.coords['two_theta']))
    )
    return DetectorData[SampleRun](da)


def mcstas_chopper_delay_from_mode(
    da: DetectorData[SampleRun],
) -> WavelengthDefinitionChopperDelay:
    mode = next(iter(d.coords['mode'] for d in da.values())).value
    if mode in ('7', '8'):
        return sc.scalar(0.00245635, unit='s') 
    if mode in ('9', '10'):
        return sc.scalar(0.0033730158730158727, unit='s')
    if mode == '16':
        return sc.scalar(0.000876984126984127, unit='s')
    raise ValueError(f'Mode {mode} is not known.')


def mcstas_modulation_period_from_mode(da: DetectorData[SampleRun]) -> ModulationPeriod:
    mode = next(iter(d.coords['mode'] for d in da.values())).value
    if mode in ('7', '8'):
        return sc.scalar(1.0 / (8 * 70), unit='s')
    if mode == '9':
        return sc.scalar(1.0 / (8 * 140), unit='s')
    if mode == '10':
        return sc.scalar(1.0 / (8 * 280), unit='s')
    if mode == '16':
        return sc.scalar(1.0 / (4 * 280), unit='s')
    raise ValueError(f'Mode {mode} is not known.')

mcstas_providers = (
    load_beer_mcstas_provider,
    mcstas_chopper_delay_from_mode,
    mcstas_modulation_period_from_mode,
)

[celinedurniak]

Another option could be to use mcstastox.
Below is an example

import mcstastox
    
with mcstastox.Read(filepath) as loaded_data:

    da = loaded_data.export_scipp(source_name="sourceMantid", sample_name="sampleMantid")['events']
    da.bins.coords["tof"] = da.bins.coords.pop("t")
    # Add variances
    # (See https://www.mcstas.org/documentation/manual/mcstas-3.5.27-manual.pdf,
    # section 2.2.1)
    da.bins.constituents['data'].variances = da.bins.constituents['data'].values**2

[jokasimr]

Can we use macstastox instead of the modified version of io.py?

[celinedurniak]

I have nothing against that. It depends on the dependencies the Scipp team wants to manage.

[jokasimr]

But I mean, does mcstastox do the same as what the modified version of io.py does?

[celinedurniak]

the loader with mcstastox does the same thing as the initial version of the loader:

• the output is a DataGroup split according to the detector banks.
• the differences are some of the coordinates, their names and the point of reference to position the components

[jokasimr]

Aha I see, so we might use it to simplify the loader but it can not replace the script you posted above.

[jokasimr]

@celinedurniak maybe you already shared all the test data with me, but in that case I've lost it, could you share it with me on nextcloud again and I'll make sure to copy it to a place where the entire data reduction team can access it?

[celinedurniak]

I put all files in the derived folder of one of the BEER's proposal: ess/data/beer/2025/287217/derived/McStasSimulations

[jokasimr]

I remember now you said that, thanks.