Dev rolf

CHAP/common/map_utils.py

@@ -6,12 +6,16 @@
 from pydantic import (
     conint,
     conlist,
+    model_validator,
     Field,
     FilePath,
 )
+from typing import Optional
 
 # Local modules
+from CHAP.common.models.common import IndexSliceConfig
 from CHAP.common.models.map import (
+    DetectorConfig,
     Detector,
     MapConfig,
 )
@@ -54,44 +58,50 @@ class MapSliceProcessor(Processor):
 
     :ivar map_config: Map configuration.
     :vartype map_config: MapConfig
-    :ivar detectors: Detector configurations.
-    :vartype detectors:
-        list[:class:`~CHAP.common.models.map.Detector`]
+    :ivar detector_config: Detector configurations.
+    :vartype detector_config: :class:`~CHAP.common.models.map.DetectorConfig`
     :ivar spec_file: SPEC file containing scan from which to read a
         slice of raw data.
-    :vartype spec_file: str
-    :ivar scan_number: Number of scan from which to read a slice of
+    :vartype spec_file: pydantic.FilePath
+    :ivar scan_numbers: Numbers of scans from which to read slices of
         raw data.
-    :vartype scan_number: int
+    :vartype scan_numbers: list[int]
+    :ivar idx_slice: Parameters for the slice of each scan to process
+        Defaults to `IndexSliceConfig()`.
+    :vartype idx_slice: CHAP.common.models.common.IndexSliceConfig, optional
     """
 
     pipeline_fields: dict = Field(
         default={
             'map_config': 'common.models.map.MapConfig',
+            'detector_config': 'common.models.map.DetectorConfig'
         },
         init_var=True)
     map_config: MapConfig
-    detectors: conlist(item_type=Detector, min_length=1)
+    detector_config: DetectorConfig
+    detectors: Optional[conlist(item_type=Detector)] = None
     spec_file: FilePath
-    scan_number: conint(gt=0)
+    scan_number: Optional[conint(gt=0)] = None
+    scan_numbers: Optional[conlist(item_type=conint(gt=0))] = None
+    idx_slice: Optional[IndexSliceConfig] = IndexSliceConfig()
 
-    def process(self, data, #spec_file, scan_number,
-                idx_slice={'start': 0, 'step': 1}):
-        """Aggregate partial spec and detector data from one scan in a
-        map, returning results in a format suitable for writing to the
-        full map container with
+    def process(self, data):
+
+        """Aggregate partial spec and detector data from one or more
+        scans in a map, returning results in a format suitable for
+        writing to the full map container with
         :class:`~CHAP.common.writer.NexusValuesWriter` or
         :class:`~CHAP.common.writer.ZarrValuesWriter`.
 
+        When all scans are adjacent in the map and `idx_slice` covers
+        each scan in full, data_points entries for the same path are
+        consolidated into a single array + slice, avoiding redundant
+        write calls.
+
         :param data: Result of `Reader.read` where at least one item
             has the value `'common.models.map.MapConfig'` for the
             `'schema'` key.
         :type data: list[PipelineData]
-        :type idx_slice: Parameters for the slice of the scan to
-            process (slice parameters are the usual for the python
-            `slice` object: `'start'`, `'stop'`, and
-            `'step'`). Defaults to `{'start': 0, 'step': '1'}`.
-        :type idx_slice: dict[str, int], optional
         :return: Slice of map data, ready to be written to a map
              container.
         :rtype: list[dict[str, Any]]
@@ -103,68 +113,223 @@ def process(self, data, #spec_file, scan_number,
         import numpy as np
 
         # Local modules
-        from chess_scanparsers import choose_scanparser
         from CHAP.common.models.map import SpecScans
 
-        ScanParser = choose_scanparser(
-            self.map_config.station, self.map_config.experiment_type)
-        scans = SpecScans(
-            spec_file=self.spec_file, scan_numbers=[self.scan_number])
-        scan = scans.get_scanparser(self.scan_number)
-
-        # Get index offset for this data slice within the map
-        npts_scan = int(scan.spec_scan_npts)
-        nscans_prev = 0
-        for scans in self.map_config.spec_scans:
-            for scan_n in scans.scan_numbers:
-                if (os.path.abspath(self.spec_file) == \
-                    os.path.abspath(self.spec_file)
-                    and scan_n == self.scan_number):
+        scans_obj = SpecScans(
+            spec_file=self.spec_file, scan_numbers=self.scan_numbers)
+        self_spec_file_abs = os.path.abspath(str(self.spec_file))
+
+        if self.map_config.experiment_type == 'EDD':
+            def get_detector_data(scan, detector, index):
+                return scan.get_detector_data(detector.get_id(), index)[0][0]
+        else:
+            def get_detector_data(scan, detector, index):
+                return scan.get_detector_data(detector.get_id(), index)
+
+        # Build flat ordered list of (abs_spec_file, scan_number) for
+        # all scans in the map to determine each scan's map position
+        map_scan_order = []
+        for spec_scans_item in self.map_config.spec_scans:
+            sf_abs = os.path.abspath(str(spec_scans_item.spec_file))
+            for sn in spec_scans_item.scan_numbers:
+                map_scan_order.append((sf_abs, sn))
+        scan_positions = {}
+        for sn in self.scan_numbers:
+            for pos, (sf, n) in enumerate(map_scan_order):
+                if sf == self_spec_file_abs and n == sn:
+                    scan_positions[sn] = pos
                     break
-                nscans_prev += 1
-        index_offset = nscans_prev * npts_scan
-
-        # Get spec scan indices to process
-        scan_indices = range(npts_scan)[slice(
-            idx_slice.get('start', 0),
-            idx_slice.get('stop', npts_scan + 1),
-            idx_slice.get('step', 1)
-        )]
-        # Get map indices to write to
-        map_indices = slice(
-            idx_slice.get('start', 0) + index_offset,
-            idx_slice.get('stop', npts_scan + 1) + index_offset,
-            idx_slice.get('step', 1)
+
+        # Process scans in map order
+        sorted_scan_numbers = sorted(
+            self.scan_numbers, key=lambda sn: scan_positions[sn])
+
+        slice_start = self.idx_slice._slice.start
+        slice_step = self.idx_slice._slice.step
+
+        # Collect per-scan metadata; assumes uniform npts across scans
+        # for index_offset calculation (index_offset = map_pos * npts)
+        per_scan = []
+        for sn in sorted_scan_numbers:
+            scan = scans_obj.get_scanparser(sn)
+            npts_scan = int(scan.spec_scan_npts)
+            index_offset = scan_positions[sn] * npts_scan
+            # Cap stop at npts_scan so map_indices and data stay in sync
+            slice_stop = min(self.idx_slice._slice.stop, npts_scan) \
+                if self.idx_slice._slice.stop > 0 else npts_scan
+            scan_indices = range(npts_scan)[
+                slice(slice_start, slice_stop, slice_step)]
+            map_indices = slice(
+                slice_start + index_offset,
+                slice_stop + index_offset,
+                slice_step,
+            )
+            per_scan.append({
+                'sn': sn, 'scan': scan, 'npts_scan': npts_scan,
+                'index_offset': index_offset,
+                'scan_indices': scan_indices,
+                'map_indices': map_indices,
+                'full': (slice_start == 0
+                         and slice_step == 1
+                         and slice_stop == npts_scan),
+            })
+
+        # Consolidate into single data_points entries when all scans
+        # are adjacent in the map and idx_slice covers each scan fully
+        sorted_positions = [scan_positions[sn] for sn in sorted_scan_numbers]
+        scans_are_adjacent = all(
+            sorted_positions[i + 1] == sorted_positions[i] + 1
+            for i in range(len(sorted_positions) - 1)
+        )
+        can_consolidate = (
+            len(per_scan) > 1
+            and scans_are_adjacent
+            and all(ps['full'] for ps in per_scan)
         )
 
-        data_points = [
-            {
-                'path': f'{self.map_config.title}/scalar_data/{s_d.label}',
-                'data': np.asarray([
-                    s_d.get_value(
-                        scans, self.scan_number, i,
-                        scalar_data=self.map_config.scalar_data)
-                    for i in scan_indices
-                ]),
-                'idx': map_indices
-            }
-            for s_d in self.map_config.all_scalar_data
-        ]
-        data_points.extend(
-            [
-                {
-                    'path': f'{self.map_config.title}/data/{det.get_id()}',
+        if can_consolidate:
+            first, last = per_scan[0], per_scan[-1]
+            merged_idx = slice(
+                first['index_offset'],
+                last['index_offset'] + last['npts_scan'],
+                1,
+            )
+            data_points = [{
+                'path': (f'{self.map_config.title}'
+                         f'/independent_dimensions/index'),
+                'data': np.arange(
+                    first['index_offset'],
+                    last['index_offset'] + last['npts_scan'],
+                ),
+                'idx': merged_idx,
+            }]
+            for s_d in self.map_config.all_scalar_data:
+                data_points.append({
+                    'path': (f'{self.map_config.title}'
+                             f'/scalar_data/{s_d.label}'),
+                    'data': np.concatenate([
+                        np.asarray([
+                            s_d.get_value(
+                                scans_obj, ps['sn'], i,
+                                scalar_data=self.map_config.scalar_data)
+                            for i in ps['scan_indices']
+                        ])
+                        for ps in per_scan
+                    ]),
+                    'idx': merged_idx,
+                })
+            for dim in self.map_config.independent_dimensions:
+                data_points.append({
+                    'path': (f'{self.map_config.title}'
+                             f'/independent_dimensions/{dim.label}'),
+                    'data': np.concatenate([
+                        np.asarray([
+                            dim.get_value(
+                                scans_obj, ps['sn'], i,
+                                scalar_data=self.map_config.scalar_data)
+                            for i in ps['scan_indices']
+                        ])
+                        for ps in per_scan
+                    ]),
+                    'idx': merged_idx,
+                })
+            for det in self.detector_config.detectors:
+                data_points.append({
+                    'path': (f'{self.map_config.title}'
+                             f'/data/{det.get_id()}'),
+                    'data': np.concatenate([
+                        np.asarray([
+                            get_detector_data(ps['scan'], det, i)
+                            for i in ps['scan_indices']
+                        ])
+                        for ps in per_scan
+                    ]),
+                    'idx': merged_idx,
+                })
+        else:
+            data_points = []
+            for ps in per_scan:
+                data_points.append({
+                    'path': (f'{self.map_config.title}'
+                             f'/independent_dimensions/index'),
+                    'data': np.asarray(
+                        [ps['index_offset'] + i for i in ps['scan_indices']]
+                    ),
+                    'idx': ps['map_indices'],
+                })
+                data_points.extend([{
+                    'path': (f'{self.map_config.title}'
+                             f'/scalar_data/{s_d.label}'),
                     'data': np.asarray([
-                        scan.get_detector_data(det.get_id(), i)
-                        for i in scan_indices
+                        s_d.get_value(
+                            scans_obj, ps['sn'], i,
+                            scalar_data=self.map_config.scalar_data)
+                        for i in ps['scan_indices']
                     ]),
-                    'idx': map_indices
-                }
-                for det in self.detectors
-            ]
-        )
+                    'idx': ps['map_indices'],
+                } for s_d in self.map_config.all_scalar_data])
+                data_points.extend([{
+                    'path': (f'{self.map_config.title}'
+                             f'/independent_dimensions/{dim.label}'),
+                    'data': np.asarray([
+                        dim.get_value(
+                            scans_obj, ps['sn'], i,
+                            scalar_data=self.map_config.scalar_data)
+                        for i in ps['scan_indices']
+                    ]),
+                    'idx': ps['map_indices'],
+                } for dim in self.map_config.independent_dimensions])
+                data_points.extend([{
+                    'path': (f'{self.map_config.title}'
+                             f'/data/{det.get_id()}'),
+                    'data': np.asarray([
+                        get_detector_data(ps['scan'], det, i)
+                        for i in ps['scan_indices']
+                    ]),
+                    'idx': ps['map_indices'],
+                } for det in self.detector_config.detectors])
         return data_points
 
+    @model_validator(mode='before')
+    @classmethod
+    def fill_scan_numbers(cls, data):
+        if not isinstance(data, dict):
+            return data
+        if 'scan_numbers' not in data or data['scan_numbers'] is None:
+            if data.get('scan_number') is not None:
+                data['scan_numbers'] = [data['scan_number']]
+        elif isinstance(data['scan_numbers'], int):
+            data['scan_numbers'] = [data['scan_numbers']]
+        elif isinstance(data['scan_numbers'], str):
+            from CHAP.utils.general import string_to_list
+            data['scan_numbers'] = string_to_list(data['scan_numbers'])
+        return data
+
+    @model_validator(mode='after')
+    def validate_scan_numbers(self):
+        if self.scan_numbers is None:
+            raise ValueError(
+                'scan_numbers is required; alternatively, provide scan_number')
+        if self.scan_number is not None \
+           and self.scan_number not in self.scan_numbers:
+            self.scan_numbers.append(self.scan_number)
+        return self
+
+    @model_validator(mode='before')
+    def fill_detector_config(cls, data):
+        if not isinstance(data, dict):
+            return data
+        if 'detector_config' not in data or data['detector_config'] is None:
+            if data.get('detectors') is not None:
+                data['detector_config'] = DetectorConfig(
+                    detectors=data['detectors']
+                )
+            else:
+                raise ValueError(
+                    'detector_config is required; alternatively, provide detectors'
+                )
+        return data
+
 
 class SpecScanToMapConfigProcessor(Processor):
     """Processor to get the

CHAP/common/models/common.py

@@ -146,6 +146,25 @@ def validate_vrange(cls, vrange, info):
                 for i in info.data['index_range']]
 
 
+class IndexSliceConfig(CHAPBaseModel):
+    """Configuration for a python `sslice` object.
+
+    :ivar start: A `start` parameter for `slice()`, defaults to 0.
+    :vartype start: int, optional
+    :ivar stop: A `stop` parameter for `slice()`, defaults to -1.
+    :vartype stop: int, optional
+    :ivar step: A `step` parameter for `slice()`, defaults to 1.
+    :vartype step: int, optional
+    """
+    start: Optional[int] = 0
+    stop: Optional[int] = -1
+    step: Optional[int] = 1
+
+    @property
+    def _slice(self):
+        return slice(self.start, self.stop, self.step)
+
+
 class UnstructuredToStructuredConfig(CHAPBaseModel):
     """Configuration class to reshape data in an
     `NXdata <https://manual.nexusformat.org/classes/base_classes/NXdata.html>`__