• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

OpenCOMPES / sed / 12147354400

03 Dec 2024 08:13PM UTC coverage: 91.773% (-0.01%) from 91.786%
12147354400

push

github

web-flow
Merge pull request #525 from OpenCOMPES/adjust_adc_ranges

Delay calibration fixes

7 of 10 new or added lines in 2 files covered. (70.0%)

1 existing line in 1 file now uncovered.

6515 of 7099 relevant lines covered (91.77%)

0.92 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

86.47
/sed/core/processor.py
1
"""This module contains the core class for the sed package
2

3
"""
4
import pathlib
1✔
5
from datetime import datetime
1✔
6
from typing import Any
1✔
7
from typing import cast
1✔
8
from typing import Dict
1✔
9
from typing import List
1✔
10
from typing import Sequence
1✔
11
from typing import Tuple
1✔
12
from typing import Union
1✔
13

14
import dask.dataframe as ddf
1✔
15
import matplotlib.pyplot as plt
1✔
16
import numpy as np
1✔
17
import pandas as pd
1✔
18
import psutil
1✔
19
import xarray as xr
1✔
20

21
from sed.binning import bin_dataframe
1✔
22
from sed.binning.binning import normalization_histogram_from_timed_dataframe
1✔
23
from sed.binning.binning import normalization_histogram_from_timestamps
1✔
24
from sed.calibrator import DelayCalibrator
1✔
25
from sed.calibrator import EnergyCalibrator
1✔
26
from sed.calibrator import MomentumCorrector
1✔
27
from sed.core.config import parse_config
1✔
28
from sed.core.config import save_config
1✔
29
from sed.core.dfops import add_time_stamped_data
1✔
30
from sed.core.dfops import apply_filter
1✔
31
from sed.core.dfops import apply_jitter
1✔
32
from sed.core.metadata import MetaHandler
1✔
33
from sed.diagnostics import grid_histogram
1✔
34
from sed.io import to_h5
1✔
35
from sed.io import to_nexus
1✔
36
from sed.io import to_tiff
1✔
37
from sed.loader import CopyTool
1✔
38
from sed.loader import get_loader
1✔
39
from sed.loader.mpes.loader import get_archiver_data
1✔
40
from sed.loader.mpes.loader import MpesLoader
1✔
41

42
N_CPU = psutil.cpu_count()
1✔
43

44

45
class SedProcessor:
1✔
46
    """Processor class of sed. Contains wrapper functions defining a work flow for data
47
    correction, calibration and binning.
48

49
    Args:
50
        metadata (dict, optional): Dict of external Metadata. Defaults to None.
51
        config (Union[dict, str], optional): Config dictionary or config file name.
52
            Defaults to None.
53
        dataframe (Union[pd.DataFrame, ddf.DataFrame], optional): dataframe to load
54
            into the class. Defaults to None.
55
        files (List[str], optional): List of files to pass to the loader defined in
56
            the config. Defaults to None.
57
        folder (str, optional): Folder containing files to pass to the loader
58
            defined in the config. Defaults to None.
59
        runs (Sequence[str], optional): List of run identifiers to pass to the loader
60
            defined in the config. Defaults to None.
61
        collect_metadata (bool): Option to collect metadata from files.
62
            Defaults to False.
63
        verbose (bool, optional): Option to print out diagnostic information.
64
            Defaults to config["core"]["verbose"] or False.
65
        **kwds: Keyword arguments passed to the reader.
66
    """
67

68
    def __init__(
1✔
69
        self,
70
        metadata: dict = None,
71
        config: Union[dict, str] = None,
72
        dataframe: Union[pd.DataFrame, ddf.DataFrame] = None,
73
        files: List[str] = None,
74
        folder: str = None,
75
        runs: Sequence[str] = None,
76
        collect_metadata: bool = False,
77
        verbose: bool = None,
78
        **kwds,
79
    ):
80
        """Processor class of sed. Contains wrapper functions defining a work flow
81
        for data correction, calibration, and binning.
82

83
        Args:
84
            metadata (dict, optional): Dict of external Metadata. Defaults to None.
85
            config (Union[dict, str], optional): Config dictionary or config file name.
86
                Defaults to None.
87
            dataframe (Union[pd.DataFrame, ddf.DataFrame], optional): dataframe to load
88
                into the class. Defaults to None.
89
            files (List[str], optional): List of files to pass to the loader defined in
90
                the config. Defaults to None.
91
            folder (str, optional): Folder containing files to pass to the loader
92
                defined in the config. Defaults to None.
93
            runs (Sequence[str], optional): List of run identifiers to pass to the loader
94
                defined in the config. Defaults to None.
95
            collect_metadata (bool, optional): Option to collect metadata from files.
96
                Defaults to False.
97
            verbose (bool, optional): Option to print out diagnostic information.
98
                Defaults to config["core"]["verbose"] or False.
99
            **kwds: Keyword arguments passed to parse_config and to the reader.
100
        """
101
        config_kwds = {
1✔
102
            key: value for key, value in kwds.items() if key in parse_config.__code__.co_varnames
103
        }
104
        for key in config_kwds.keys():
1✔
105
            del kwds[key]
1✔
106
        self._config = parse_config(config, **config_kwds)
1✔
107
        num_cores = self._config.get("binning", {}).get("num_cores", N_CPU - 1)
1✔
108
        if num_cores >= N_CPU:
1✔
109
            num_cores = N_CPU - 1
1✔
110
        self._config["binning"]["num_cores"] = num_cores
1✔
111

112
        if verbose is None:
1✔
113
            self.verbose = self._config["core"].get("verbose", False)
1✔
114
        else:
115
            self.verbose = verbose
1✔
116

117
        self._dataframe: Union[pd.DataFrame, ddf.DataFrame] = None
1✔
118
        self._timed_dataframe: Union[pd.DataFrame, ddf.DataFrame] = None
1✔
119
        self._files: List[str] = []
1✔
120

121
        self._binned: xr.DataArray = None
1✔
122
        self._pre_binned: xr.DataArray = None
1✔
123
        self._normalization_histogram: xr.DataArray = None
1✔
124
        self._normalized: xr.DataArray = None
1✔
125

126
        self._attributes = MetaHandler(meta=metadata)
1✔
127

128
        loader_name = self._config["core"]["loader"]
1✔
129
        self.loader = get_loader(
1✔
130
            loader_name=loader_name,
131
            config=self._config,
132
        )
133

134
        self.ec = EnergyCalibrator(
1✔
135
            loader=get_loader(
136
                loader_name=loader_name,
137
                config=self._config,
138
            ),
139
            config=self._config,
140
        )
141

142
        self.mc = MomentumCorrector(
1✔
143
            config=self._config,
144
        )
145

146
        self.dc = DelayCalibrator(
1✔
147
            config=self._config,
148
        )
149

150
        self.use_copy_tool = self._config.get("core", {}).get(
1✔
151
            "use_copy_tool",
152
            False,
153
        )
154
        if self.use_copy_tool:
1✔
155
            try:
1✔
156
                self.ct = CopyTool(
1✔
157
                    source=self._config["core"]["copy_tool_source"],
158
                    dest=self._config["core"]["copy_tool_dest"],
159
                    **self._config["core"].get("copy_tool_kwds", {}),
160
                )
161
            except KeyError:
1✔
162
                self.use_copy_tool = False
1✔
163

164
        # Load data if provided:
165
        if dataframe is not None or files is not None or folder is not None or runs is not None:
1✔
166
            self.load(
1✔
167
                dataframe=dataframe,
168
                metadata=metadata,
169
                files=files,
170
                folder=folder,
171
                runs=runs,
172
                collect_metadata=collect_metadata,
173
                **kwds,
174
            )
175

176
    def __repr__(self):
1✔
177
        if self._dataframe is None:
1✔
178
            df_str = "Dataframe: No Data loaded"
1✔
179
        else:
180
            df_str = self._dataframe.__repr__()
1✔
181
        pretty_str = df_str + "\n" + "Metadata: " + "\n" + self._attributes.__repr__()
1✔
182
        return pretty_str
1✔
183

184
    def _repr_html_(self):
1✔
185
        html = "<div>"
×
186

187
        if self._dataframe is None:
×
188
            df_html = "Dataframe: No Data loaded"
×
189
        else:
190
            df_html = self._dataframe._repr_html_()
×
191

192
        html += f"<details><summary>Dataframe</summary>{df_html}</details>"
×
193

194
        # Add expandable section for attributes
195
        html += "<details><summary>Metadata</summary>"
×
196
        html += "<div style='padding-left: 10px;'>"
×
197
        html += self._attributes._repr_html_()
×
198
        html += "</div></details>"
×
199

200
        html += "</div>"
×
201

202
        return html
×
203

204
    ## Suggestion:
205
    # @property
206
    # def overview_panel(self):
207
    #     """Provides an overview panel with plots of different data attributes."""
208
    #     self.view_event_histogram(dfpid=2, backend="matplotlib")
209

210
    @property
1✔
211
    def dataframe(self) -> Union[pd.DataFrame, ddf.DataFrame]:
1✔
212
        """Accessor to the underlying dataframe.
213

214
        Returns:
215
            Union[pd.DataFrame, ddf.DataFrame]: Dataframe object.
216
        """
217
        return self._dataframe
1✔
218

219
    @dataframe.setter
1✔
220
    def dataframe(self, dataframe: Union[pd.DataFrame, ddf.DataFrame]):
1✔
221
        """Setter for the underlying dataframe.
222

223
        Args:
224
            dataframe (Union[pd.DataFrame, ddf.DataFrame]): The dataframe object to set.
225
        """
226
        if not isinstance(dataframe, (pd.DataFrame, ddf.DataFrame)) or not isinstance(
1✔
227
            dataframe,
228
            self._dataframe.__class__,
229
        ):
230
            raise ValueError(
1✔
231
                "'dataframe' has to be a Pandas or Dask dataframe and has to be of the same kind "
232
                "as the dataframe loaded into the SedProcessor!.\n"
233
                f"Loaded type: {self._dataframe.__class__}, provided type: {dataframe}.",
234
            )
235
        self._dataframe = dataframe
1✔
236

237
    @property
1✔
238
    def timed_dataframe(self) -> Union[pd.DataFrame, ddf.DataFrame]:
1✔
239
        """Accessor to the underlying timed_dataframe.
240

241
        Returns:
242
            Union[pd.DataFrame, ddf.DataFrame]: Timed Dataframe object.
243
        """
244
        return self._timed_dataframe
1✔
245

246
    @timed_dataframe.setter
1✔
247
    def timed_dataframe(self, timed_dataframe: Union[pd.DataFrame, ddf.DataFrame]):
1✔
248
        """Setter for the underlying timed dataframe.
249

250
        Args:
251
            timed_dataframe (Union[pd.DataFrame, ddf.DataFrame]): The timed dataframe object to set
252
        """
253
        if not isinstance(timed_dataframe, (pd.DataFrame, ddf.DataFrame)) or not isinstance(
×
254
            timed_dataframe,
255
            self._timed_dataframe.__class__,
256
        ):
257
            raise ValueError(
×
258
                "'timed_dataframe' has to be a Pandas or Dask dataframe and has to be of the same "
259
                "kind as the dataframe loaded into the SedProcessor!.\n"
260
                f"Loaded type: {self._timed_dataframe.__class__}, "
261
                f"provided type: {timed_dataframe}.",
262
            )
263
        self._timed_dataframe = timed_dataframe
×
264

265
    @property
1✔
266
    def attributes(self) -> MetaHandler:
1✔
267
        """Accessor to the metadata dict.
268

269
        Returns:
270
            MetaHandler: The metadata object
271
        """
272
        return self._attributes
1✔
273

274
    def add_attribute(self, attributes: dict, name: str, **kwds):
1✔
275
        """Function to add element to the attributes dict.
276

277
        Args:
278
            attributes (dict): The attributes dictionary object to add.
279
            name (str): Key under which to add the dictionary to the attributes.
280
        """
281
        self._attributes.add(
1✔
282
            entry=attributes,
283
            name=name,
284
            **kwds,
285
        )
286

287
    @property
1✔
288
    def config(self) -> Dict[Any, Any]:
1✔
289
        """Getter attribute for the config dictionary
290

291
        Returns:
292
            Dict: The config dictionary.
293
        """
294
        return self._config
1✔
295

296
    @property
1✔
297
    def files(self) -> List[str]:
1✔
298
        """Getter attribute for the list of files
299

300
        Returns:
301
            List[str]: The list of loaded files
302
        """
303
        return self._files
1✔
304

305
    @property
1✔
306
    def binned(self) -> xr.DataArray:
1✔
307
        """Getter attribute for the binned data array
308

309
        Returns:
310
            xr.DataArray: The binned data array
311
        """
312
        if self._binned is None:
1✔
313
            raise ValueError("No binned data available, need to compute histogram first!")
×
314
        return self._binned
1✔
315

316
    @property
1✔
317
    def normalized(self) -> xr.DataArray:
1✔
318
        """Getter attribute for the normalized data array
319

320
        Returns:
321
            xr.DataArray: The normalized data array
322
        """
323
        if self._normalized is None:
1✔
324
            raise ValueError(
×
325
                "No normalized data available, compute data with normalization enabled!",
326
            )
327
        return self._normalized
1✔
328

329
    @property
1✔
330
    def normalization_histogram(self) -> xr.DataArray:
1✔
331
        """Getter attribute for the normalization histogram
332

333
        Returns:
334
            xr.DataArray: The normalization histogram
335
        """
336
        if self._normalization_histogram is None:
1✔
337
            raise ValueError("No normalization histogram available, generate histogram first!")
×
338
        return self._normalization_histogram
1✔
339

340
    def cpy(self, path: Union[str, List[str]]) -> Union[str, List[str]]:
1✔
341
        """Function to mirror a list of files or a folder from a network drive to a
342
        local storage. Returns either the original or the copied path to the given
343
        path. The option to use this functionality is set by
344
        config["core"]["use_copy_tool"].
345

346
        Args:
347
            path (Union[str, List[str]]): Source path or path list.
348

349
        Returns:
350
            Union[str, List[str]]: Source or destination path or path list.
351
        """
352
        if self.use_copy_tool:
1✔
353
            if isinstance(path, list):
1✔
354
                path_out = []
1✔
355
                for file in path:
1✔
356
                    path_out.append(self.ct.copy(file))
1✔
357
                return path_out
1✔
358

359
            return self.ct.copy(path)
×
360

361
        if isinstance(path, list):
1✔
362
            return path
1✔
363

364
        return path
1✔
365

366
    def load(
1✔
367
        self,
368
        dataframe: Union[pd.DataFrame, ddf.DataFrame] = None,
369
        metadata: dict = None,
370
        files: List[str] = None,
371
        folder: str = None,
372
        runs: Sequence[str] = None,
373
        collect_metadata: bool = False,
374
        **kwds,
375
    ):
376
        """Load tabular data of single events into the dataframe object in the class.
377

378
        Args:
379
            dataframe (Union[pd.DataFrame, ddf.DataFrame], optional): data in tabular
380
                format. Accepts anything which can be interpreted by pd.DataFrame as
381
                an input. Defaults to None.
382
            metadata (dict, optional): Dict of external Metadata. Defaults to None.
383
            files (List[str], optional): List of file paths to pass to the loader.
384
                Defaults to None.
385
            runs (Sequence[str], optional): List of run identifiers to pass to the
386
                loader. Defaults to None.
387
            folder (str, optional): Folder path to pass to the loader.
388
                Defaults to None.
389
            collect_metadata (bool, optional): Option for collecting metadata in the reader.
390
            **kwds: Keyword parameters passed to the reader.
391

392
        Raises:
393
            ValueError: Raised if no valid input is provided.
394
        """
395
        if metadata is None:
1✔
396
            metadata = {}
1✔
397
        if dataframe is not None:
1✔
398
            timed_dataframe = kwds.pop("timed_dataframe", None)
1✔
399
        elif runs is not None:
1✔
400
            # If runs are provided, we only use the copy tool if also folder is provided.
401
            # In that case, we copy the whole provided base folder tree, and pass the copied
402
            # version to the loader as base folder to look for the runs.
403
            if folder is not None:
1✔
404
                dataframe, timed_dataframe, metadata = self.loader.read_dataframe(
1✔
405
                    folders=cast(str, self.cpy(folder)),
406
                    runs=runs,
407
                    metadata=metadata,
408
                    collect_metadata=collect_metadata,
409
                    **kwds,
410
                )
411
            else:
412
                dataframe, timed_dataframe, metadata = self.loader.read_dataframe(
×
413
                    runs=runs,
414
                    metadata=metadata,
415
                    collect_metadata=collect_metadata,
416
                    **kwds,
417
                )
418

419
        elif folder is not None:
1✔
420
            dataframe, timed_dataframe, metadata = self.loader.read_dataframe(
1✔
421
                folders=cast(str, self.cpy(folder)),
422
                metadata=metadata,
423
                collect_metadata=collect_metadata,
424
                **kwds,
425
            )
426
        elif files is not None:
1✔
427
            dataframe, timed_dataframe, metadata = self.loader.read_dataframe(
1✔
428
                files=cast(List[str], self.cpy(files)),
429
                metadata=metadata,
430
                collect_metadata=collect_metadata,
431
                **kwds,
432
            )
433
        else:
434
            raise ValueError(
1✔
435
                "Either 'dataframe', 'files', 'folder', or 'runs' needs to be provided!",
436
            )
437

438
        self._dataframe = dataframe
1✔
439
        self._timed_dataframe = timed_dataframe
1✔
440
        self._files = self.loader.files
1✔
441

442
        for key in metadata:
1✔
443
            self._attributes.add(
1✔
444
                entry=metadata[key],
445
                name=key,
446
                duplicate_policy="merge",
447
            )
448

449
    def filter_column(
1✔
450
        self,
451
        column: str,
452
        min_value: float = -np.inf,
453
        max_value: float = np.inf,
454
    ) -> None:
455
        """Filter values in a column which are outside of a given range
456

457
        Args:
458
            column (str): Name of the column to filter
459
            min_value (float, optional): Minimum value to keep. Defaults to None.
460
            max_value (float, optional): Maximum value to keep. Defaults to None.
461
        """
462
        if column != "index" and column not in self._dataframe.columns:
1✔
463
            raise KeyError(f"Column {column} not found in dataframe!")
1✔
464
        if min_value >= max_value:
1✔
465
            raise ValueError("min_value has to be smaller than max_value!")
1✔
466
        if self._dataframe is not None:
1✔
467
            self._dataframe = apply_filter(
1✔
468
                self._dataframe,
469
                col=column,
470
                lower_bound=min_value,
471
                upper_bound=max_value,
472
            )
473
        if self._timed_dataframe is not None and column in self._timed_dataframe.columns:
1✔
474
            self._timed_dataframe = apply_filter(
1✔
475
                self._timed_dataframe,
476
                column,
477
                lower_bound=min_value,
478
                upper_bound=max_value,
479
            )
480
        metadata = {
1✔
481
            "filter": {
482
                "column": column,
483
                "min_value": min_value,
484
                "max_value": max_value,
485
            },
486
        }
487
        self._attributes.add(metadata, "filter", duplicate_policy="merge")
1✔
488

489
    # Momentum calibration workflow
490
    # 1. Bin raw detector data for distortion correction
491
    def bin_and_load_momentum_calibration(
1✔
492
        self,
493
        df_partitions: Union[int, Sequence[int]] = 100,
494
        axes: List[str] = None,
495
        bins: List[int] = None,
496
        ranges: Sequence[Tuple[float, float]] = None,
497
        plane: int = 0,
498
        width: int = 5,
499
        apply: bool = False,
500
        **kwds,
501
    ):
502
        """1st step of momentum correction work flow. Function to do an initial binning
503
        of the dataframe loaded to the class, slice a plane from it using an
504
        interactive view, and load it into the momentum corrector class.
505

506
        Args:
507
            df_partitions (Union[int, Sequence[int]], optional): Number of dataframe partitions
508
                to use for the initial binning. Defaults to 100.
509
            axes (List[str], optional): Axes to bin.
510
                Defaults to config["momentum"]["axes"].
511
            bins (List[int], optional): Bin numbers to use for binning.
512
                Defaults to config["momentum"]["bins"].
513
            ranges (List[Tuple], optional): Ranges to use for binning.
514
                Defaults to config["momentum"]["ranges"].
515
            plane (int, optional): Initial value for the plane slider. Defaults to 0.
516
            width (int, optional): Initial value for the width slider. Defaults to 5.
517
            apply (bool, optional): Option to directly apply the values and select the
518
                slice. Defaults to False.
519
            **kwds: Keyword argument passed to the pre_binning function.
520
        """
521
        self._pre_binned = self.pre_binning(
1✔
522
            df_partitions=df_partitions,
523
            axes=axes,
524
            bins=bins,
525
            ranges=ranges,
526
            **kwds,
527
        )
528

529
        self.mc.load_data(data=self._pre_binned)
1✔
530
        self.mc.select_slicer(plane=plane, width=width, apply=apply)
1✔
531

532
    # 2. Generate the spline warp correction from momentum features.
533
    # Either autoselect features, or input features from view above.
534
    def define_features(
1✔
535
        self,
536
        features: np.ndarray = None,
537
        rotation_symmetry: int = 6,
538
        auto_detect: bool = False,
539
        include_center: bool = True,
540
        apply: bool = False,
541
        **kwds,
542
    ):
543
        """2. Step of the distortion correction workflow: Define feature points in
544
        momentum space. They can be either manually selected using a GUI tool, be
545
        provided as list of feature points, or auto-generated using a
546
        feature-detection algorithm.
547

548
        Args:
549
            features (np.ndarray, optional): np.ndarray of features. Defaults to None.
550
            rotation_symmetry (int, optional): Number of rotational symmetry axes.
551
                Defaults to 6.
552
            auto_detect (bool, optional): Whether to auto-detect the features.
553
                Defaults to False.
554
            include_center (bool, optional): Option to include a point at the center
555
                in the feature list. Defaults to True.
556
            apply (bool, optional): Option to directly apply the values and select the
557
                slice. Defaults to False.
558
            **kwds: Keyword arguments for ``MomentumCorrector.feature_extract()`` and
559
                ``MomentumCorrector.feature_select()``.
560
        """
561
        if auto_detect:  # automatic feature selection
1✔
562
            sigma = kwds.pop("sigma", self._config["momentum"]["sigma"])
×
563
            fwhm = kwds.pop("fwhm", self._config["momentum"]["fwhm"])
×
564
            sigma_radius = kwds.pop(
×
565
                "sigma_radius",
566
                self._config["momentum"]["sigma_radius"],
567
            )
568
            self.mc.feature_extract(
×
569
                sigma=sigma,
570
                fwhm=fwhm,
571
                sigma_radius=sigma_radius,
572
                rotsym=rotation_symmetry,
573
                **kwds,
574
            )
575
            features = self.mc.peaks
×
576

577
        self.mc.feature_select(
1✔
578
            rotsym=rotation_symmetry,
579
            include_center=include_center,
580
            features=features,
581
            apply=apply,
582
            **kwds,
583
        )
584

585
    # 3. Generate the spline warp correction from momentum features.
586
    # If no features have been selected before, use class defaults.
587
    def generate_splinewarp(
1✔
588
        self,
589
        use_center: bool = None,
590
        verbose: bool = None,
591
        **kwds,
592
    ):
593
        """3. Step of the distortion correction workflow: Generate the correction
594
        function restoring the symmetry in the image using a splinewarp algorithm.
595

596
        Args:
597
            use_center (bool, optional): Option to use the position of the
598
                center point in the correction. Default is read from config, or set to True.
599
            verbose (bool, optional): Option to print out diagnostic information.
600
                Defaults to config["core"]["verbose"].
601
            **kwds: Keyword arguments for MomentumCorrector.spline_warp_estimate().
602
        """
603
        if verbose is None:
1✔
604
            verbose = self.verbose
1✔
605

606
        self.mc.spline_warp_estimate(use_center=use_center, verbose=verbose, **kwds)
1✔
607

608
        if self.mc.slice is not None and verbose:
1✔
609
            print("Original slice with reference features")
1✔
610
            self.mc.view(annotated=True, backend="bokeh", crosshair=True)
1✔
611

612
            print("Corrected slice with target features")
1✔
613
            self.mc.view(
1✔
614
                image=self.mc.slice_corrected,
615
                annotated=True,
616
                points={"feats": self.mc.ptargs},
617
                backend="bokeh",
618
                crosshair=True,
619
            )
620

621
            print("Original slice with target features")
1✔
622
            self.mc.view(
1✔
623
                image=self.mc.slice,
624
                points={"feats": self.mc.ptargs},
625
                annotated=True,
626
                backend="bokeh",
627
            )
628

629
    # 3a. Save spline-warp parameters to config file.
630
    def save_splinewarp(
1✔
631
        self,
632
        filename: str = None,
633
        overwrite: bool = False,
634
    ):
635
        """Save the generated spline-warp parameters to the folder config file.
636

637
        Args:
638
            filename (str, optional): Filename of the config dictionary to save to.
639
                Defaults to "sed_config.yaml" in the current folder.
640
            overwrite (bool, optional): Option to overwrite the present dictionary.
641
                Defaults to False.
642
        """
643
        if filename is None:
1✔
644
            filename = "sed_config.yaml"
×
645
        if len(self.mc.correction) == 0:
1✔
646
            raise ValueError("No momentum correction parameters to save!")
×
647
        correction = {}
1✔
648
        for key, value in self.mc.correction.items():
1✔
649
            if key in ["reference_points", "target_points", "cdeform_field", "rdeform_field"]:
1✔
650
                continue
1✔
651
            if key in ["use_center", "rotation_symmetry"]:
1✔
652
                correction[key] = value
1✔
653
            elif key in ["center_point", "ascale"]:
1✔
654
                correction[key] = [float(i) for i in value]
1✔
655
            elif key in ["outer_points", "feature_points"]:
1✔
656
                correction[key] = []
1✔
657
                for point in value:
1✔
658
                    correction[key].append([float(i) for i in point])
1✔
659
            else:
660
                correction[key] = float(value)
1✔
661

662
        if "creation_date" not in correction:
1✔
663
            correction["creation_date"] = datetime.now().timestamp()
×
664

665
        config = {
1✔
666
            "momentum": {
667
                "correction": correction,
668
            },
669
        }
670
        save_config(config, filename, overwrite)
1✔
671
        print(f'Saved momentum correction parameters to "{filename}".')
1✔
672

673
    # 4. Pose corrections. Provide interactive interface for correcting
674
    # scaling, shift and rotation
675
    def pose_adjustment(
1✔
676
        self,
677
        transformations: Dict[str, Any] = None,
678
        apply: bool = False,
679
        use_correction: bool = True,
680
        reset: bool = True,
681
        verbose: bool = None,
682
        **kwds,
683
    ):
684
        """3. step of the distortion correction workflow: Generate an interactive panel
685
        to adjust affine transformations that are applied to the image. Applies first
686
        a scaling, next an x/y translation, and last a rotation around the center of
687
        the image.
688

689
        Args:
690
            transformations (dict, optional): Dictionary with transformations.
691
                Defaults to self.transformations or config["momentum"]["transformations"].
692
            apply (bool, optional): Option to directly apply the provided
693
                transformations. Defaults to False.
694
            use_correction (bool, option): Whether to use the spline warp correction
695
                or not. Defaults to True.
696
            reset (bool, optional): Option to reset the correction before transformation.
697
                Defaults to True.
698
            verbose (bool, optional): Option to print out diagnostic information.
699
                Defaults to config["core"]["verbose"].
700
            **kwds: Keyword parameters defining defaults for the transformations:
701

702
                - **scale** (float): Initial value of the scaling slider.
703
                - **xtrans** (float): Initial value of the xtrans slider.
704
                - **ytrans** (float): Initial value of the ytrans slider.
705
                - **angle** (float): Initial value of the angle slider.
706
        """
707
        if verbose is None:
1✔
708
            verbose = self.verbose
1✔
709

710
        # Generate homography as default if no distortion correction has been applied
711
        if self.mc.slice_corrected is None:
1✔
712
            if self.mc.slice is None:
1✔
713
                self.mc.slice = np.zeros(self._config["momentum"]["bins"][0:2])
1✔
714
                self.mc.bin_ranges = self._config["momentum"]["ranges"]
1✔
715
            self.mc.slice_corrected = self.mc.slice
1✔
716

717
        if not use_correction:
1✔
718
            self.mc.reset_deformation()
1✔
719

720
        if self.mc.cdeform_field is None or self.mc.rdeform_field is None:
1✔
721
            # Generate distortion correction from config values
722
            self.mc.spline_warp_estimate(verbose=verbose)
1✔
723

724
        self.mc.pose_adjustment(
1✔
725
            transformations=transformations,
726
            apply=apply,
727
            reset=reset,
728
            verbose=verbose,
729
            **kwds,
730
        )
731

732
    # 4a. Save pose adjustment parameters to config file.
733
    def save_transformations(
1✔
734
        self,
735
        filename: str = None,
736
        overwrite: bool = False,
737
    ):
738
        """Save the pose adjustment parameters to the folder config file.
739

740
        Args:
741
            filename (str, optional): Filename of the config dictionary to save to.
742
                Defaults to "sed_config.yaml" in the current folder.
743
            overwrite (bool, optional): Option to overwrite the present dictionary.
744
                Defaults to False.
745
        """
746
        if filename is None:
1✔
747
            filename = "sed_config.yaml"
×
748
        if len(self.mc.transformations) == 0:
1✔
749
            raise ValueError("No momentum transformation parameters to save!")
×
750
        transformations = {}
1✔
751
        for key, value in self.mc.transformations.items():
1✔
752
            transformations[key] = float(value)
1✔
753

754
        if "creation_date" not in transformations:
1✔
755
            transformations["creation_date"] = datetime.now().timestamp()
×
756

757
        config = {
1✔
758
            "momentum": {
759
                "transformations": transformations,
760
            },
761
        }
762
        save_config(config, filename, overwrite)
1✔
763
        print(f'Saved momentum transformation parameters to "{filename}".')
1✔
764

765
    # 5. Apply the momentum correction to the dataframe
766
    def apply_momentum_correction(
1✔
767
        self,
768
        preview: bool = False,
769
        verbose: bool = None,
770
        **kwds,
771
    ):
772
        """Applies the distortion correction and pose adjustment (optional)
773
        to the dataframe.
774

775
        Args:
776
            preview (bool, optional): Option to preview the first elements of the data frame.
777
                Defaults to False.
778
            verbose (bool, optional): Option to print out diagnostic information.
779
                Defaults to config["core"]["verbose"].
780
            **kwds: Keyword parameters for ``MomentumCorrector.apply_correction``:
781

782
                - **rdeform_field** (np.ndarray, optional): Row deformation field.
783
                - **cdeform_field** (np.ndarray, optional): Column deformation field.
784
                - **inv_dfield** (np.ndarray, optional): Inverse deformation field.
785

786
        """
787
        if verbose is None:
1✔
788
            verbose = self.verbose
1✔
789

790
        x_column = self._config["dataframe"]["x_column"]
1✔
791
        y_column = self._config["dataframe"]["y_column"]
1✔
792

793
        if self._dataframe is not None:
1✔
794
            if verbose:
1✔
795
                print("Adding corrected X/Y columns to dataframe:")
1✔
796
            df, metadata = self.mc.apply_corrections(
1✔
797
                df=self._dataframe,
798
                verbose=verbose,
799
                **kwds,
800
            )
801
            if (
1✔
802
                self._timed_dataframe is not None
803
                and x_column in self._timed_dataframe.columns
804
                and y_column in self._timed_dataframe.columns
805
            ):
806
                tdf, _ = self.mc.apply_corrections(
1✔
807
                    self._timed_dataframe,
808
                    verbose=False,
809
                    **kwds,
810
                )
811

812
            # Add Metadata
813
            self._attributes.add(
1✔
814
                metadata,
815
                "momentum_correction",
816
                duplicate_policy="merge",
817
            )
818
            self._dataframe = df
1✔
819
            if (
1✔
820
                self._timed_dataframe is not None
821
                and x_column in self._timed_dataframe.columns
822
                and y_column in self._timed_dataframe.columns
823
            ):
824
                self._timed_dataframe = tdf
1✔
825
        else:
826
            raise ValueError("No dataframe loaded!")
×
827
        if preview:
1✔
828
            print(self._dataframe.head(10))
×
829
        else:
830
            if self.verbose:
1✔
831
                print(self._dataframe)
1✔
832

833
    # Momentum calibration work flow
834
    # 1. Calculate momentum calibration
835
    def calibrate_momentum_axes(
1✔
836
        self,
837
        point_a: Union[np.ndarray, List[int]] = None,
838
        point_b: Union[np.ndarray, List[int]] = None,
839
        k_distance: float = None,
840
        k_coord_a: Union[np.ndarray, List[float]] = None,
841
        k_coord_b: Union[np.ndarray, List[float]] = np.array([0.0, 0.0]),
842
        equiscale: bool = True,
843
        apply=False,
844
    ):
845
        """1. step of the momentum calibration workflow. Calibrate momentum
846
        axes using either provided pixel coordinates of a high-symmetry point and its
847
        distance to the BZ center, or the k-coordinates of two points in the BZ
848
        (depending on the equiscale option). Opens an interactive panel for selecting
849
        the points.
850

851
        Args:
852
            point_a (Union[np.ndarray, List[int]]): Pixel coordinates of the first
853
                point used for momentum calibration.
854
            point_b (Union[np.ndarray, List[int]], optional): Pixel coordinates of the
855
                second point used for momentum calibration.
856
                Defaults to config["momentum"]["center_pixel"].
857
            k_distance (float, optional): Momentum distance between point a and b.
858
                Needs to be provided if no specific k-coordinates for the two points
859
                are given. Defaults to None.
860
            k_coord_a (Union[np.ndarray, List[float]], optional): Momentum coordinate
861
                of the first point used for calibration. Used if equiscale is False.
862
                Defaults to None.
863
            k_coord_b (Union[np.ndarray, List[float]], optional): Momentum coordinate
864
                of the second point used for calibration. Defaults to [0.0, 0.0].
865
            equiscale (bool, optional): Option to apply different scales to kx and ky.
866
                If True, the distance between points a and b, and the absolute
867
                position of point a are used for defining the scale. If False, the
868
                scale is calculated from the k-positions of both points a and b.
869
                Defaults to True.
870
            apply (bool, optional): Option to directly store the momentum calibration
871
                in the class. Defaults to False.
872
        """
873
        if point_b is None:
1✔
874
            point_b = self._config["momentum"]["center_pixel"]
1✔
875

876
        self.mc.select_k_range(
1✔
877
            point_a=point_a,
878
            point_b=point_b,
879
            k_distance=k_distance,
880
            k_coord_a=k_coord_a,
881
            k_coord_b=k_coord_b,
882
            equiscale=equiscale,
883
            apply=apply,
884
        )
885

886
    # 1a. Save momentum calibration parameters to config file.
887
    def save_momentum_calibration(
1✔
888
        self,
889
        filename: str = None,
890
        overwrite: bool = False,
891
    ):
892
        """Save the generated momentum calibration parameters to the folder config file.
893

894
        Args:
895
            filename (str, optional): Filename of the config dictionary to save to.
896
                Defaults to "sed_config.yaml" in the current folder.
897
            overwrite (bool, optional): Option to overwrite the present dictionary.
898
                Defaults to False.
899
        """
900
        if filename is None:
1✔
901
            filename = "sed_config.yaml"
×
902
        if len(self.mc.calibration) == 0:
1✔
903
            raise ValueError("No momentum calibration parameters to save!")
×
904
        calibration = {}
1✔
905
        for key, value in self.mc.calibration.items():
1✔
906
            if key in ["kx_axis", "ky_axis", "grid", "extent"]:
1✔
907
                continue
1✔
908

909
            calibration[key] = float(value)
1✔
910

911
        if "creation_date" not in calibration:
1✔
912
            calibration["creation_date"] = datetime.now().timestamp()
×
913

914
        config = {"momentum": {"calibration": calibration}}
1✔
915
        save_config(config, filename, overwrite)
1✔
916
        print(f"Saved momentum calibration parameters to {filename}")
1✔
917

918
    # 2. Apply correction and calibration to the dataframe
919
    def apply_momentum_calibration(
1✔
920
        self,
921
        calibration: dict = None,
922
        preview: bool = False,
923
        verbose: bool = None,
924
        **kwds,
925
    ):
926
        """2. step of the momentum calibration work flow: Apply the momentum
927
        calibration stored in the class to the dataframe. If corrected X/Y axis exist,
928
        these are used.
929

930
        Args:
931
            calibration (dict, optional): Optional dictionary with calibration data to
932
                use. Defaults to None.
933
            preview (bool, optional): Option to preview the first elements of the data frame.
934
                Defaults to False.
935
            verbose (bool, optional): Option to print out diagnostic information.
936
                Defaults to config["core"]["verbose"].
937
            **kwds: Keyword args passed to ``DelayCalibrator.append_delay_axis``.
938
        """
939
        if verbose is None:
1✔
940
            verbose = self.verbose
1✔
941

942
        x_column = self._config["dataframe"]["x_column"]
1✔
943
        y_column = self._config["dataframe"]["y_column"]
1✔
944

945
        if self._dataframe is not None:
1✔
946
            if verbose:
1✔
947
                print("Adding kx/ky columns to dataframe:")
1✔
948
            df, metadata = self.mc.append_k_axis(
1✔
949
                df=self._dataframe,
950
                calibration=calibration,
951
                **kwds,
952
            )
953
            if (
1✔
954
                self._timed_dataframe is not None
955
                and x_column in self._timed_dataframe.columns
956
                and y_column in self._timed_dataframe.columns
957
            ):
958
                tdf, _ = self.mc.append_k_axis(
1✔
959
                    df=self._timed_dataframe,
960
                    calibration=calibration,
961
                    **kwds,
962
                )
963

964
            # Add Metadata
965
            self._attributes.add(
1✔
966
                metadata,
967
                "momentum_calibration",
968
                duplicate_policy="merge",
969
            )
970
            self._dataframe = df
1✔
971
            if (
1✔
972
                self._timed_dataframe is not None
973
                and x_column in self._timed_dataframe.columns
974
                and y_column in self._timed_dataframe.columns
975
            ):
976
                self._timed_dataframe = tdf
1✔
977
        else:
978
            raise ValueError("No dataframe loaded!")
×
979
        if preview:
1✔
980
            print(self._dataframe.head(10))
×
981
        else:
982
            if self.verbose:
1✔
983
                print(self._dataframe)
1✔
984

985
    # Energy correction workflow
986
    # 1. Adjust the energy correction parameters
987
    def adjust_energy_correction(
1✔
988
        self,
989
        correction_type: str = None,
990
        amplitude: float = None,
991
        center: Tuple[float, float] = None,
992
        apply=False,
993
        **kwds,
994
    ):
995
        """1. step of the energy correction workflow: Opens an interactive plot to
996
        adjust the parameters for the TOF/energy correction. Also pre-bins the data if
997
        they are not present yet.
998

999
        Args:
1000
            correction_type (str, optional): Type of correction to apply to the TOF
1001
                axis. Valid values are:
1002

1003
                - 'spherical'
1004
                - 'Lorentzian'
1005
                - 'Gaussian'
1006
                - 'Lorentzian_asymmetric'
1007

1008
                Defaults to config["energy"]["correction_type"].
1009
            amplitude (float, optional): Amplitude of the correction.
1010
                Defaults to config["energy"]["correction"]["amplitude"].
1011
            center (Tuple[float, float], optional): Center X/Y coordinates for the
1012
                correction. Defaults to config["energy"]["correction"]["center"].
1013
            apply (bool, optional): Option to directly apply the provided or default
1014
                correction parameters. Defaults to False.
1015
            **kwds: Keyword parameters passed to ``EnergyCalibrator.adjust_energy_correction()``.
1016
        """
1017
        if self._pre_binned is None:
1✔
1018
            print(
1✔
1019
                "Pre-binned data not present, binning using defaults from config...",
1020
            )
1021
            self._pre_binned = self.pre_binning()
1✔
1022

1023
        self.ec.adjust_energy_correction(
1✔
1024
            self._pre_binned,
1025
            correction_type=correction_type,
1026
            amplitude=amplitude,
1027
            center=center,
1028
            apply=apply,
1029
            **kwds,
1030
        )
1031

1032
    # 1a. Save energy correction parameters to config file.
1033
    def save_energy_correction(
1✔
1034
        self,
1035
        filename: str = None,
1036
        overwrite: bool = False,
1037
    ):
1038
        """Save the generated energy correction parameters to the folder config file.
1039

1040
        Args:
1041
            filename (str, optional): Filename of the config dictionary to save to.
1042
                Defaults to "sed_config.yaml" in the current folder.
1043
            overwrite (bool, optional): Option to overwrite the present dictionary.
1044
                Defaults to False.
1045
        """
1046
        if filename is None:
1✔
1047
            filename = "sed_config.yaml"
1✔
1048
        if len(self.ec.correction) == 0:
1✔
1049
            raise ValueError("No energy correction parameters to save!")
×
1050
        correction = {}
1✔
1051
        for key, val in self.ec.correction.items():
1✔
1052
            if key == "correction_type":
1✔
1053
                correction[key] = val
1✔
1054
            elif key == "center":
1✔
1055
                correction[key] = [float(i) for i in val]
1✔
1056
            else:
1057
                correction[key] = float(val)
1✔
1058

1059
        if "creation_date" not in correction:
1✔
1060
            correction["creation_date"] = datetime.now().timestamp()
×
1061

1062
        config = {"energy": {"correction": correction}}
1✔
1063
        save_config(config, filename, overwrite)
1✔
1064
        print(f"Saved energy correction parameters to {filename}")
1✔
1065

1066
    # 2. Apply energy correction to dataframe
1067
    def apply_energy_correction(
1✔
1068
        self,
1069
        correction: dict = None,
1070
        preview: bool = False,
1071
        verbose: bool = None,
1072
        **kwds,
1073
    ):
1074
        """2. step of the energy correction workflow: Apply the energy correction
1075
        parameters stored in the class to the dataframe.
1076

1077
        Args:
1078
            correction (dict, optional): Dictionary containing the correction
1079
                parameters. Defaults to config["energy"]["calibration"].
1080
            preview (bool, optional): Option to preview the first elements of the data frame.
1081
                Defaults to False.
1082
            verbose (bool, optional): Option to print out diagnostic information.
1083
                Defaults to config["core"]["verbose"].
1084
            **kwds:
1085
                Keyword args passed to ``EnergyCalibrator.apply_energy_correction()``.
1086
        """
1087
        if verbose is None:
1✔
1088
            verbose = self.verbose
1✔
1089

1090
        tof_column = self._config["dataframe"]["tof_column"]
1✔
1091

1092
        if self._dataframe is not None:
1✔
1093
            if verbose:
1✔
1094
                print("Applying energy correction to dataframe...")
1✔
1095
            df, metadata = self.ec.apply_energy_correction(
1✔
1096
                df=self._dataframe,
1097
                correction=correction,
1098
                verbose=verbose,
1099
                **kwds,
1100
            )
1101
            if self._timed_dataframe is not None and tof_column in self._timed_dataframe.columns:
1✔
1102
                tdf, _ = self.ec.apply_energy_correction(
1✔
1103
                    df=self._timed_dataframe,
1104
                    correction=correction,
1105
                    verbose=False,
1106
                    **kwds,
1107
                )
1108

1109
            # Add Metadata
1110
            self._attributes.add(
1✔
1111
                metadata,
1112
                "energy_correction",
1113
            )
1114
            self._dataframe = df
1✔
1115
            if self._timed_dataframe is not None and tof_column in self._timed_dataframe.columns:
1✔
1116
                self._timed_dataframe = tdf
1✔
1117
        else:
1118
            raise ValueError("No dataframe loaded!")
×
1119
        if preview:
1✔
1120
            print(self._dataframe.head(10))
×
1121
        else:
1122
            if verbose:
1✔
1123
                print(self._dataframe)
×
1124

1125
    # Energy calibrator workflow
1126
    # 1. Load and normalize data
1127
    def load_bias_series(
1✔
1128
        self,
1129
        binned_data: Union[xr.DataArray, Tuple[np.ndarray, np.ndarray, np.ndarray]] = None,
1130
        data_files: List[str] = None,
1131
        axes: List[str] = None,
1132
        bins: List = None,
1133
        ranges: Sequence[Tuple[float, float]] = None,
1134
        biases: np.ndarray = None,
1135
        bias_key: str = None,
1136
        normalize: bool = None,
1137
        span: int = None,
1138
        order: int = None,
1139
    ):
1140
        """1. step of the energy calibration workflow: Load and bin data from
1141
        single-event files, or load binned bias/TOF traces.
1142

1143
        Args:
1144
            binned_data (Union[xr.DataArray, Tuple[np.ndarray, np.ndarray, np.ndarray]], optional):
1145
                Binned data If provided as DataArray, Needs to contain dimensions
1146
                config["dataframe"]["tof_column"] and config["dataframe"]["bias_column"]. If
1147
                provided as tuple, needs to contain elements tof, biases, traces.
1148
            data_files (List[str], optional): list of file paths to bin
1149
            axes (List[str], optional): bin axes.
1150
                Defaults to config["dataframe"]["tof_column"].
1151
            bins (List, optional): number of bins.
1152
                Defaults to config["energy"]["bins"].
1153
            ranges (Sequence[Tuple[float, float]], optional): bin ranges.
1154
                Defaults to config["energy"]["ranges"].
1155
            biases (np.ndarray, optional): Bias voltages used. If missing, bias
1156
                voltages are extracted from the data files.
1157
            bias_key (str, optional): hdf5 path where bias values are stored.
1158
                Defaults to config["energy"]["bias_key"].
1159
            normalize (bool, optional): Option to normalize traces.
1160
                Defaults to config["energy"]["normalize"].
1161
            span (int, optional): span smoothing parameters of the LOESS method
1162
                (see ``scipy.signal.savgol_filter()``).
1163
                Defaults to config["energy"]["normalize_span"].
1164
            order (int, optional): order smoothing parameters of the LOESS method
1165
                (see ``scipy.signal.savgol_filter()``).
1166
                Defaults to config["energy"]["normalize_order"].
1167
        """
1168
        if binned_data is not None:
1✔
1169
            if isinstance(binned_data, xr.DataArray):
1✔
1170
                if (
1✔
1171
                    self._config["dataframe"]["tof_column"] not in binned_data.dims
1172
                    or self._config["dataframe"]["bias_column"] not in binned_data.dims
1173
                ):
1174
                    raise ValueError(
1✔
1175
                        "If binned_data is provided as an xarray, it needs to contain dimensions "
1176
                        f"'{self._config['dataframe']['tof_column']}' and "
1177
                        f"'{self._config['dataframe']['bias_column']}'!.",
1178
                    )
1179
                tof = binned_data.coords[self._config["dataframe"]["tof_column"]].values
1✔
1180
                biases = binned_data.coords[self._config["dataframe"]["bias_column"]].values
1✔
1181
                traces = binned_data.values[:, :]
1✔
1182
            else:
1183
                try:
1✔
1184
                    (tof, biases, traces) = binned_data
1✔
1185
                except ValueError as exc:
1✔
1186
                    raise ValueError(
1✔
1187
                        "If binned_data is provided as tuple, it needs to contain "
1188
                        "(tof, biases, traces)!",
1189
                    ) from exc
1190
            self.ec.load_data(biases=biases, traces=traces, tof=tof)
1✔
1191

1192
        elif data_files is not None:
1✔
1193
            self.ec.bin_data(
1✔
1194
                data_files=cast(List[str], self.cpy(data_files)),
1195
                axes=axes,
1196
                bins=bins,
1197
                ranges=ranges,
1198
                biases=biases,
1199
                bias_key=bias_key,
1200
            )
1201

1202
        else:
1203
            raise ValueError("Either binned_data or data_files needs to be provided!")
1✔
1204

1205
        if (normalize is not None and normalize is True) or (
1✔
1206
            normalize is None and self._config["energy"]["normalize"]
1207
        ):
1208
            if span is None:
1✔
1209
                span = self._config["energy"]["normalize_span"]
1✔
1210
            if order is None:
1✔
1211
                order = self._config["energy"]["normalize_order"]
1✔
1212
            self.ec.normalize(smooth=True, span=span, order=order)
1✔
1213
        self.ec.view(
1✔
1214
            traces=self.ec.traces_normed,
1215
            xaxis=self.ec.tof,
1216
            backend="bokeh",
1217
        )
1218

1219
    # 2. extract ranges and get peak positions
1220
    def find_bias_peaks(
1✔
1221
        self,
1222
        ranges: Union[List[Tuple], Tuple],
1223
        ref_id: int = 0,
1224
        infer_others: bool = True,
1225
        mode: str = "replace",
1226
        radius: int = None,
1227
        peak_window: int = None,
1228
        apply: bool = False,
1229
    ):
1230
        """2. step of the energy calibration workflow: Find a peak within a given range
1231
        for the indicated reference trace, and tries to find the same peak for all
1232
        other traces. Uses fast_dtw to align curves, which might not be too good if the
1233
        shape of curves changes qualitatively. Ideally, choose a reference trace in the
1234
        middle of the set, and don't choose the range too narrow around the peak.
1235
        Alternatively, a list of ranges for all traces can be provided.
1236

1237
        Args:
1238
            ranges (Union[List[Tuple], Tuple]): Tuple of TOF values indicating a range.
1239
                Alternatively, a list of ranges for all traces can be given.
1240
            ref_id (int, optional): The id of the trace the range refers to.
1241
                Defaults to 0.
1242
            infer_others (bool, optional): Whether to determine the range for the other
1243
                traces. Defaults to True.
1244
            mode (str, optional): Whether to "add" or "replace" existing ranges.
1245
                Defaults to "replace".
1246
            radius (int, optional): Radius parameter for fast_dtw.
1247
                Defaults to config["energy"]["fastdtw_radius"].
1248
            peak_window (int, optional): Peak_window parameter for the peak detection
1249
                algorithm. amount of points that have to have to behave monotonously
1250
                around a peak. Defaults to config["energy"]["peak_window"].
1251
            apply (bool, optional): Option to directly apply the provided parameters.
1252
                Defaults to False.
1253
        """
1254
        if radius is None:
1✔
1255
            radius = self._config["energy"]["fastdtw_radius"]
1✔
1256
        if peak_window is None:
1✔
1257
            peak_window = self._config["energy"]["peak_window"]
1✔
1258
        if not infer_others:
1✔
1259
            self.ec.add_ranges(
1✔
1260
                ranges=ranges,
1261
                ref_id=ref_id,
1262
                infer_others=infer_others,
1263
                mode=mode,
1264
                radius=radius,
1265
            )
1266
            print(self.ec.featranges)
1✔
1267
            try:
1✔
1268
                self.ec.feature_extract(peak_window=peak_window)
1✔
1269
                self.ec.view(
1✔
1270
                    traces=self.ec.traces_normed,
1271
                    segs=self.ec.featranges,
1272
                    xaxis=self.ec.tof,
1273
                    peaks=self.ec.peaks,
1274
                    backend="bokeh",
1275
                )
1276
            except IndexError:
×
1277
                print("Could not determine all peaks!")
×
1278
                raise
×
1279
        else:
1280
            # New adjustment tool
1281
            assert isinstance(ranges, tuple)
1✔
1282
            self.ec.adjust_ranges(
1✔
1283
                ranges=ranges,
1284
                ref_id=ref_id,
1285
                traces=self.ec.traces_normed,
1286
                infer_others=infer_others,
1287
                radius=radius,
1288
                peak_window=peak_window,
1289
                apply=apply,
1290
            )
1291

1292
    # 3. Fit the energy calibration relation
1293
    def calibrate_energy_axis(
1✔
1294
        self,
1295
        ref_id: int,
1296
        ref_energy: float,
1297
        method: str = None,
1298
        energy_scale: str = None,
1299
        verbose: bool = None,
1300
        **kwds,
1301
    ):
1302
        """3. Step of the energy calibration workflow: Calculate the calibration
1303
        function for the energy axis, and apply it to the dataframe. Two
1304
        approximations are implemented, a (normally 3rd order) polynomial
1305
        approximation, and a d^2/(t-t0)^2 relation.
1306

1307
        Args:
1308
            ref_id (int): id of the trace at the bias where the reference energy is
1309
                given.
1310
            ref_energy (float): Absolute energy of the detected feature at the bias
1311
                of ref_id
1312
            method (str, optional): Method for determining the energy calibration.
1313

1314
                - **'lmfit'**: Energy calibration using lmfit and 1/t^2 form.
1315
                - **'lstsq'**, **'lsqr'**: Energy calibration using polynomial form.
1316

1317
                Defaults to config["energy"]["calibration_method"]
1318
            energy_scale (str, optional): Direction of increasing energy scale.
1319

1320
                - **'kinetic'**: increasing energy with decreasing TOF.
1321
                - **'binding'**: increasing energy with increasing TOF.
1322

1323
                Defaults to config["energy"]["energy_scale"]
1324
            verbose (bool, optional): Option to print out diagnostic information.
1325
                Defaults to config["core"]["verbose"].
1326
            **kwds**: Keyword parameters passed to ``EnergyCalibrator.calibrate()``.
1327
        """
1328
        if verbose is None:
1✔
1329
            verbose = self.verbose
1✔
1330

1331
        if method is None:
1✔
1332
            method = self._config["energy"]["calibration_method"]
1✔
1333

1334
        if energy_scale is None:
1✔
1335
            energy_scale = self._config["energy"]["energy_scale"]
1✔
1336

1337
        self.ec.calibrate(
1✔
1338
            ref_id=ref_id,
1339
            ref_energy=ref_energy,
1340
            method=method,
1341
            energy_scale=energy_scale,
1342
            verbose=verbose,
1343
            **kwds,
1344
        )
1345
        if verbose:
1✔
1346
            if self.ec.traces_normed is not None:
1✔
1347
                print("Quality of Calibration:")
1✔
1348
                self.ec.view(
1✔
1349
                    traces=self.ec.traces_normed,
1350
                    xaxis=self.ec.calibration["axis"],
1351
                    align=True,
1352
                    energy_scale=energy_scale,
1353
                    backend="bokeh",
1354
                )
1355
            print("E/TOF relationship:")
1✔
1356
            self.ec.view(
1✔
1357
                traces=self.ec.calibration["axis"][None, :],
1358
                xaxis=self.ec.tof,
1359
                backend="matplotlib",
1360
                show_legend=False,
1361
            )
1362
            if energy_scale == "kinetic":
1✔
1363
                plt.scatter(
1✔
1364
                    self.ec.peaks[:, 0],
1365
                    -(self.ec.biases - self.ec.biases[ref_id]) + ref_energy,
1366
                    s=50,
1367
                    c="k",
1368
                )
1369
            elif energy_scale == "binding":
1✔
1370
                plt.scatter(
1✔
1371
                    self.ec.peaks[:, 0],
1372
                    self.ec.biases - self.ec.biases[ref_id] + ref_energy,
1373
                    s=50,
1374
                    c="k",
1375
                )
1376
            else:
1377
                raise ValueError(
×
1378
                    'energy_scale needs to be either "binding" or "kinetic"',
1379
                    f", got {energy_scale}.",
1380
                )
1381
            plt.xlabel("Time-of-flight", fontsize=15)
1✔
1382
            plt.ylabel("Energy (eV)", fontsize=15)
1✔
1383
            plt.show()
1✔
1384

1385
    # 3a. Save energy calibration parameters to config file.
1386
    def save_energy_calibration(
1✔
1387
        self,
1388
        filename: str = None,
1389
        overwrite: bool = False,
1390
    ):
1391
        """Save the generated energy calibration parameters to the folder config file.
1392

1393
        Args:
1394
            filename (str, optional): Filename of the config dictionary to save to.
1395
                Defaults to "sed_config.yaml" in the current folder.
1396
            overwrite (bool, optional): Option to overwrite the present dictionary.
1397
                Defaults to False.
1398
        """
1399
        if filename is None:
1✔
1400
            filename = "sed_config.yaml"
×
1401
        if len(self.ec.calibration) == 0:
1✔
1402
            raise ValueError("No energy calibration parameters to save!")
×
1403
        calibration = {}
1✔
1404
        for key, value in self.ec.calibration.items():
1✔
1405
            if key in ["axis", "refid", "Tmat", "bvec"]:
1✔
1406
                continue
1✔
1407
            if key == "energy_scale":
1✔
1408
                calibration[key] = value
1✔
1409
            elif key == "coeffs":
1✔
1410
                calibration[key] = [float(i) for i in value]
1✔
1411
            else:
1412
                calibration[key] = float(value)
1✔
1413

1414
        if "creation_date" not in calibration:
1✔
1415
            calibration["creation_date"] = datetime.now().timestamp()
×
1416

1417
        config = {"energy": {"calibration": calibration}}
1✔
1418
        save_config(config, filename, overwrite)
1✔
1419
        print(f'Saved energy calibration parameters to "{filename}".')
1✔
1420

1421
    # 4. Apply energy calibration to the dataframe
1422
    def append_energy_axis(
1✔
1423
        self,
1424
        calibration: dict = None,
1425
        preview: bool = False,
1426
        verbose: bool = None,
1427
        **kwds,
1428
    ):
1429
        """4. step of the energy calibration workflow: Apply the calibration function
1430
        to to the dataframe. Two approximations are implemented, a (normally 3rd order)
1431
        polynomial approximation, and a d^2/(t-t0)^2 relation. a calibration dictionary
1432
        can be provided.
1433

1434
        Args:
1435
            calibration (dict, optional): Calibration dict containing calibration
1436
                parameters. Overrides calibration from class or config.
1437
                Defaults to None.
1438
            preview (bool): Option to preview the first elements of the data frame.
1439
            verbose (bool, optional): Option to print out diagnostic information.
1440
                Defaults to config["core"]["verbose"].
1441
            **kwds:
1442
                Keyword args passed to ``EnergyCalibrator.append_energy_axis()``.
1443
        """
1444
        if verbose is None:
1✔
1445
            verbose = self.verbose
1✔
1446

1447
        tof_column = self._config["dataframe"]["tof_column"]
1✔
1448

1449
        if self._dataframe is not None:
1✔
1450
            if verbose:
1✔
1451
                print("Adding energy column to dataframe:")
1✔
1452
            df, metadata = self.ec.append_energy_axis(
1✔
1453
                df=self._dataframe,
1454
                calibration=calibration,
1455
                verbose=verbose,
1456
                **kwds,
1457
            )
1458
            if self._timed_dataframe is not None and tof_column in self._timed_dataframe.columns:
1✔
1459
                tdf, _ = self.ec.append_energy_axis(
1✔
1460
                    df=self._timed_dataframe,
1461
                    calibration=calibration,
1462
                    verbose=False,
1463
                    **kwds,
1464
                )
1465

1466
            # Add Metadata
1467
            self._attributes.add(
1✔
1468
                metadata,
1469
                "energy_calibration",
1470
                duplicate_policy="merge",
1471
            )
1472
            self._dataframe = df
1✔
1473
            if self._timed_dataframe is not None and tof_column in self._timed_dataframe.columns:
1✔
1474
                self._timed_dataframe = tdf
1✔
1475

1476
        else:
1477
            raise ValueError("No dataframe loaded!")
×
1478
        if preview:
1✔
1479
            print(self._dataframe.head(10))
×
1480
        else:
1481
            if verbose:
1✔
1482
                print(self._dataframe)
1✔
1483

1484
    def add_energy_offset(
1✔
1485
        self,
1486
        constant: float = None,
1487
        columns: Union[str, Sequence[str]] = None,
1488
        weights: Union[float, Sequence[float]] = None,
1489
        reductions: Union[str, Sequence[str]] = None,
1490
        preserve_mean: Union[bool, Sequence[bool]] = None,
1491
        preview: bool = False,
1492
        verbose: bool = None,
1493
    ) -> None:
1494
        """Shift the energy axis of the dataframe by a given amount.
1495

1496
        Args:
1497
            constant (float, optional): The constant to shift the energy axis by.
1498
            columns (Union[str, Sequence[str]]): Name of the column(s) to apply the shift from.
1499
            weights (Union[float, Sequence[float]]): weights to apply to the columns.
1500
                Can also be used to flip the sign (e.g. -1). Defaults to 1.
1501
            preserve_mean (bool): Whether to subtract the mean of the column before applying the
1502
                shift. Defaults to False.
1503
            reductions (str): The reduction to apply to the column. Should be an available method
1504
                of dask.dataframe.Series. For example "mean". In this case the function is applied
1505
                to the column to generate a single value for the whole dataset. If None, the shift
1506
                is applied per-dataframe-row. Defaults to None. Currently only "mean" is supported.
1507
            preview (bool, optional): Option to preview the first elements of the data frame.
1508
                Defaults to False.
1509
            verbose (bool, optional): Option to print out diagnostic information.
1510
                Defaults to config["core"]["verbose"].
1511

1512
        Raises:
1513
            ValueError: If the energy column is not in the dataframe.
1514
        """
1515
        if verbose is None:
1✔
1516
            verbose = self.verbose
1✔
1517

1518
        energy_column = self._config["dataframe"]["energy_column"]
1✔
1519
        if energy_column not in self._dataframe.columns:
1✔
1520
            raise ValueError(
1✔
1521
                f"Energy column {energy_column} not found in dataframe! "
1522
                "Run `append_energy_axis()` first.",
1523
            )
1524
        if self.dataframe is not None:
1✔
1525
            if verbose:
1✔
1526
                print("Adding energy offset to dataframe:")
1✔
1527
            df, metadata = self.ec.add_offsets(
1✔
1528
                df=self._dataframe,
1529
                constant=constant,
1530
                columns=columns,
1531
                energy_column=energy_column,
1532
                weights=weights,
1533
                reductions=reductions,
1534
                preserve_mean=preserve_mean,
1535
                verbose=verbose,
1536
            )
1537
            if self._timed_dataframe is not None and energy_column in self._timed_dataframe.columns:
1✔
1538
                tdf, _ = self.ec.add_offsets(
1✔
1539
                    df=self._timed_dataframe,
1540
                    constant=constant,
1541
                    columns=columns,
1542
                    energy_column=energy_column,
1543
                    weights=weights,
1544
                    reductions=reductions,
1545
                    preserve_mean=preserve_mean,
1546
                )
1547

1548
            self._attributes.add(
1✔
1549
                metadata,
1550
                "add_energy_offset",
1551
                # TODO: allow only appending when no offset along this column(s) was applied
1552
                # TODO: clear memory of modifications if the energy axis is recalculated
1553
                duplicate_policy="append",
1554
            )
1555
            self._dataframe = df
1✔
1556
            if self._timed_dataframe is not None and energy_column in self._timed_dataframe.columns:
1✔
1557
                self._timed_dataframe = tdf
1✔
1558
        else:
1559
            raise ValueError("No dataframe loaded!")
×
1560
        if preview:
1✔
1561
            print(self._dataframe.head(10))
×
1562
        elif verbose:
1✔
1563
            print(self._dataframe)
1✔
1564

1565
    def save_energy_offset(
1✔
1566
        self,
1567
        filename: str = None,
1568
        overwrite: bool = False,
1569
    ):
1570
        """Save the generated energy calibration parameters to the folder config file.
1571

1572
        Args:
1573
            filename (str, optional): Filename of the config dictionary to save to.
1574
                Defaults to "sed_config.yaml" in the current folder.
1575
            overwrite (bool, optional): Option to overwrite the present dictionary.
1576
                Defaults to False.
1577
        """
1578
        if filename is None:
×
1579
            filename = "sed_config.yaml"
×
1580
        if len(self.ec.offsets) == 0:
×
1581
            raise ValueError("No energy offset parameters to save!")
×
1582

1583
        if "creation_date" not in self.ec.offsets.keys():
×
1584
            self.ec.offsets["creation_date"] = datetime.now().timestamp()
×
1585

1586
        config = {"energy": {"offsets": self.ec.offsets}}
×
1587
        save_config(config, filename, overwrite)
×
1588
        print(f'Saved energy offset parameters to "{filename}".')
×
1589

1590
    def append_tof_ns_axis(
1✔
1591
        self,
1592
        preview: bool = False,
1593
        verbose: bool = None,
1594
        **kwds,
1595
    ):
1596
        """Convert time-of-flight channel steps to nanoseconds.
1597

1598
        Args:
1599
            tof_ns_column (str, optional): Name of the generated column containing the
1600
                time-of-flight in nanosecond.
1601
                Defaults to config["dataframe"]["tof_ns_column"].
1602
            preview (bool, optional): Option to preview the first elements of the data frame.
1603
                Defaults to False.
1604
            verbose (bool, optional): Option to print out diagnostic information.
1605
                Defaults to config["core"]["verbose"].
1606
            **kwds: additional arguments are passed to ``EnergyCalibrator.tof_step_to_ns()``.
1607

1608
        """
1609
        if verbose is None:
1✔
1610
            verbose = self.verbose
1✔
1611

1612
        tof_column = self._config["dataframe"]["tof_column"]
1✔
1613

1614
        if self._dataframe is not None:
1✔
1615
            if verbose:
1✔
1616
                print("Adding time-of-flight column in nanoseconds to dataframe:")
1✔
1617
            # TODO assert order of execution through metadata
1618

1619
            df, metadata = self.ec.append_tof_ns_axis(
1✔
1620
                df=self._dataframe,
1621
                **kwds,
1622
            )
1623
            if self._timed_dataframe is not None and tof_column in self._timed_dataframe.columns:
1✔
1624
                tdf, _ = self.ec.append_tof_ns_axis(
1✔
1625
                    df=self._timed_dataframe,
1626
                    **kwds,
1627
                )
1628

1629
            self._attributes.add(
1✔
1630
                metadata,
1631
                "tof_ns_conversion",
1632
                duplicate_policy="overwrite",
1633
            )
1634
            self._dataframe = df
1✔
1635
            if self._timed_dataframe is not None and tof_column in self._timed_dataframe.columns:
1✔
1636
                self._timed_dataframe = tdf
1✔
1637
        else:
1638
            raise ValueError("No dataframe loaded!")
×
1639
        if preview:
1✔
1640
            print(self._dataframe.head(10))
×
1641
        else:
1642
            if verbose:
1✔
1643
                print(self._dataframe)
1✔
1644

1645
    def align_dld_sectors(
1✔
1646
        self,
1647
        sector_delays: np.ndarray = None,
1648
        preview: bool = False,
1649
        verbose: bool = None,
1650
        **kwds,
1651
    ):
1652
        """Align the 8s sectors of the HEXTOF endstation.
1653

1654
        Args:
1655
            sector_delays (np.ndarray, optional): Array containing the sector delays. Defaults to
1656
                config["dataframe"]["sector_delays"].
1657
            preview (bool, optional): Option to preview the first elements of the data frame.
1658
                Defaults to False.
1659
            verbose (bool, optional): Option to print out diagnostic information.
1660
                Defaults to config["core"]["verbose"].
1661
            **kwds: additional arguments are passed to ``EnergyCalibrator.align_dld_sectors()``.
1662
        """
1663
        if verbose is None:
1✔
1664
            verbose = self.verbose
1✔
1665

1666
        tof_column = self._config["dataframe"]["tof_column"]
1✔
1667

1668
        if self._dataframe is not None:
1✔
1669
            if verbose:
1✔
1670
                print("Aligning 8s sectors of dataframe")
1✔
1671
            # TODO assert order of execution through metadata
1672

1673
            df, metadata = self.ec.align_dld_sectors(
1✔
1674
                df=self._dataframe,
1675
                sector_delays=sector_delays,
1676
                **kwds,
1677
            )
1678
            if self._timed_dataframe is not None and tof_column in self._timed_dataframe.columns:
1✔
1679
                tdf, _ = self.ec.align_dld_sectors(
×
1680
                    df=self._timed_dataframe,
1681
                    sector_delays=sector_delays,
1682
                    **kwds,
1683
                )
1684

1685
            self._attributes.add(
1✔
1686
                metadata,
1687
                "dld_sector_alignment",
1688
                duplicate_policy="raise",
1689
            )
1690
            self._dataframe = df
1✔
1691
            if self._timed_dataframe is not None and tof_column in self._timed_dataframe.columns:
1✔
1692
                self._timed_dataframe = tdf
×
1693
        else:
1694
            raise ValueError("No dataframe loaded!")
×
1695
        if preview:
1✔
1696
            print(self._dataframe.head(10))
×
1697
        else:
1698
            if verbose:
1✔
1699
                print(self._dataframe)
1✔
1700

1701
    # Delay calibration function
1702
    def calibrate_delay_axis(
1✔
1703
        self,
1704
        delay_range: Tuple[float, float] = None,
1705
        read_delay_ranges: bool = True,
1706
        datafile: str = None,
1707
        preview: bool = False,
1708
        verbose: bool = None,
1709
        **kwds,
1710
    ):
1711
        """Append delay column to dataframe. Either provide delay ranges, or read
1712
        them from a file.
1713

1714
        Args:
1715
            delay_range (Tuple[float, float], optional): The scanned delay range in
1716
                picoseconds. Defaults to None.
1717
            read_delay_ranges (bool, optional): Option whether to read the delay ranges from the
1718
                data. Defaults to True. If false, parameters in the config will be used.
1719
            datafile (str, optional): The file from which to read the delay ranges.
1720
                Defaults to the first file of the dataset.
1721
            preview (bool, optional): Option to preview the first elements of the data frame.
1722
                Defaults to False.
1723
            verbose (bool, optional): Option to print out diagnostic information.
1724
                Defaults to config["core"]["verbose"].
1725
            **kwds: Keyword args passed to ``DelayCalibrator.append_delay_axis``.
1726
        """
1727
        if verbose is None:
1✔
1728
            verbose = self.verbose
1✔
1729

1730
        adc_column = self._config["dataframe"]["adc_column"]
1✔
1731
        if adc_column not in self._dataframe.columns:
1✔
1732
            raise ValueError(f"ADC column {adc_column} not found in dataframe, cannot calibrate!")
×
1733

1734
        if self._dataframe is not None:
1✔
1735
            if verbose:
1✔
1736
                print("Adding delay column to dataframe:")
1✔
1737

1738
            if read_delay_ranges:
1✔
1739
                try:
1✔
1740
                    datafile = self._files[0]
1✔
NEW
1741
                except IndexError:
×
NEW
1742
                    print(
×
1743
                        "No datafile available, specify either",
1744
                        " 'datafile' or 'delay_range'",
1745
                    )
NEW
1746
                    raise
×
1747

1748
            df, metadata = self.dc.append_delay_axis(
1✔
1749
                self._dataframe,
1750
                delay_range=delay_range,
1751
                datafile=datafile,
1752
                verbose=verbose,
1753
                **kwds,
1754
            )
1755
            if self._timed_dataframe is not None and adc_column in self._timed_dataframe.columns:
1✔
1756
                tdf, _ = self.dc.append_delay_axis(
1✔
1757
                    self._timed_dataframe,
1758
                    delay_range=delay_range,
1759
                    datafile=datafile,
1760
                    verbose=False,
1761
                    **kwds,
1762
                )
1763

1764
            # Add Metadata
1765
            self._attributes.add(
1✔
1766
                metadata,
1767
                "delay_calibration",
1768
                duplicate_policy="overwrite",
1769
            )
1770
            self._dataframe = df
1✔
1771
            if self._timed_dataframe is not None and adc_column in self._timed_dataframe.columns:
1✔
1772
                self._timed_dataframe = tdf
1✔
1773
        else:
1774
            raise ValueError("No dataframe loaded!")
×
1775
        if preview:
1✔
UNCOV
1776
            print(self._dataframe.head(10))
×
1777
        else:
1778
            if self.verbose:
1✔
1779
                print(self._dataframe)
1✔
1780

1781
    def save_delay_calibration(
1✔
1782
        self,
1783
        filename: str = None,
1784
        overwrite: bool = False,
1785
    ) -> None:
1786
        """Save the generated delay calibration parameters to the folder config file.
1787

1788
        Args:
1789
            filename (str, optional): Filename of the config dictionary to save to.
1790
                Defaults to "sed_config.yaml" in the current folder.
1791
            overwrite (bool, optional): Option to overwrite the present dictionary.
1792
                Defaults to False.
1793
        """
1794
        if filename is None:
1✔
1795
            filename = "sed_config.yaml"
×
1796

1797
        if len(self.dc.calibration) == 0:
1✔
1798
            raise ValueError("No delay calibration parameters to save!")
×
1799
        calibration = {}
1✔
1800
        for key, value in self.dc.calibration.items():
1✔
1801
            if key == "datafile":
1✔
1802
                calibration[key] = value
1✔
1803
            elif key in ["adc_range", "delay_range", "delay_range_mm"]:
1✔
1804
                calibration[key] = [float(i) for i in value]
1✔
1805
            else:
1806
                calibration[key] = float(value)
1✔
1807

1808
        if "creation_date" not in calibration:
1✔
1809
            calibration["creation_date"] = datetime.now().timestamp()
×
1810

1811
        config = {
1✔
1812
            "delay": {
1813
                "calibration": calibration,
1814
            },
1815
        }
1816
        save_config(config, filename, overwrite)
1✔
1817

1818
    def add_delay_offset(
1✔
1819
        self,
1820
        constant: float = None,
1821
        flip_delay_axis: bool = None,
1822
        columns: Union[str, Sequence[str]] = None,
1823
        weights: Union[float, Sequence[float]] = 1.0,
1824
        reductions: Union[str, Sequence[str]] = None,
1825
        preserve_mean: Union[bool, Sequence[bool]] = False,
1826
        preview: bool = False,
1827
        verbose: bool = None,
1828
    ) -> None:
1829
        """Shift the delay axis of the dataframe by a constant or other columns.
1830

1831
        Args:
1832
            constant (float, optional): The constant to shift the delay axis by.
1833
            flip_delay_axis (bool, optional): Option to reverse the direction of the delay axis.
1834
            columns (Union[str, Sequence[str]]): Name of the column(s) to apply the shift from.
1835
            weights (Union[float, Sequence[float]]): weights to apply to the columns.
1836
                Can also be used to flip the sign (e.g. -1). Defaults to 1.
1837
            preserve_mean (bool): Whether to subtract the mean of the column before applying the
1838
                shift. Defaults to False.
1839
            reductions (str): The reduction to apply to the column. Should be an available method
1840
                of dask.dataframe.Series. For example "mean". In this case the function is applied
1841
                to the column to generate a single value for the whole dataset. If None, the shift
1842
                is applied per-dataframe-row. Defaults to None. Currently only "mean" is supported.
1843
            preview (bool, optional): Option to preview the first elements of the data frame.
1844
                Defaults to False.
1845
            verbose (bool, optional): Option to print out diagnostic information.
1846
                Defaults to config["core"]["verbose"].
1847

1848
        Raises:
1849
            ValueError: If the delay column is not in the dataframe.
1850
        """
1851
        if verbose is None:
1✔
1852
            verbose = self.verbose
1✔
1853

1854
        delay_column = self._config["dataframe"]["delay_column"]
1✔
1855
        if delay_column not in self._dataframe.columns:
1✔
1856
            raise ValueError(f"Delay column {delay_column} not found in dataframe! ")
1✔
1857

1858
        if self.dataframe is not None:
1✔
1859
            if verbose:
1✔
1860
                print("Adding delay offset to dataframe:")
1✔
1861
            df, metadata = self.dc.add_offsets(
1✔
1862
                df=self._dataframe,
1863
                constant=constant,
1864
                flip_delay_axis=flip_delay_axis,
1865
                columns=columns,
1866
                delay_column=delay_column,
1867
                weights=weights,
1868
                reductions=reductions,
1869
                preserve_mean=preserve_mean,
1870
                verbose=verbose,
1871
            )
1872
            if self._timed_dataframe is not None and delay_column in self._timed_dataframe.columns:
1✔
1873
                tdf, _ = self.dc.add_offsets(
1✔
1874
                    df=self._timed_dataframe,
1875
                    constant=constant,
1876
                    flip_delay_axis=flip_delay_axis,
1877
                    columns=columns,
1878
                    delay_column=delay_column,
1879
                    weights=weights,
1880
                    reductions=reductions,
1881
                    preserve_mean=preserve_mean,
1882
                    verbose=False,
1883
                )
1884

1885
            self._attributes.add(
1✔
1886
                metadata,
1887
                "delay_offset",
1888
                duplicate_policy="append",
1889
            )
1890
            self._dataframe = df
1✔
1891
            if self._timed_dataframe is not None and delay_column in self._timed_dataframe.columns:
1✔
1892
                self._timed_dataframe = tdf
1✔
1893
        else:
1894
            raise ValueError("No dataframe loaded!")
×
1895
        if preview:
1✔
1896
            print(self._dataframe.head(10))
1✔
1897
        else:
1898
            if verbose:
1✔
1899
                print(self._dataframe)
1✔
1900

1901
    def save_delay_offsets(
1✔
1902
        self,
1903
        filename: str = None,
1904
        overwrite: bool = False,
1905
    ) -> None:
1906
        """Save the generated delay calibration parameters to the folder config file.
1907

1908
        Args:
1909
            filename (str, optional): Filename of the config dictionary to save to.
1910
                Defaults to "sed_config.yaml" in the current folder.
1911
            overwrite (bool, optional): Option to overwrite the present dictionary.
1912
                Defaults to False.
1913
        """
1914
        if filename is None:
1✔
1915
            filename = "sed_config.yaml"
×
1916
        if len(self.dc.offsets) == 0:
1✔
1917
            raise ValueError("No delay offset parameters to save!")
×
1918

1919
        if "creation_date" not in self.ec.offsets.keys():
1✔
1920
            self.ec.offsets["creation_date"] = datetime.now().timestamp()
1✔
1921

1922
        config = {
1✔
1923
            "delay": {
1924
                "offsets": self.dc.offsets,
1925
            },
1926
        }
1927
        save_config(config, filename, overwrite)
1✔
1928
        print(f'Saved delay offset parameters to "{filename}".')
1✔
1929

1930
    def save_workflow_params(
1✔
1931
        self,
1932
        filename: str = None,
1933
        overwrite: bool = False,
1934
    ) -> None:
1935
        """run all save calibration parameter methods
1936

1937
        Args:
1938
            filename (str, optional): Filename of the config dictionary to save to.
1939
                Defaults to "sed_config.yaml" in the current folder.
1940
            overwrite (bool, optional): Option to overwrite the present dictionary.
1941
                Defaults to False.
1942
        """
1943
        for method in [
×
1944
            self.save_splinewarp,
1945
            self.save_transformations,
1946
            self.save_momentum_calibration,
1947
            self.save_energy_correction,
1948
            self.save_energy_calibration,
1949
            self.save_energy_offset,
1950
            self.save_delay_calibration,
1951
            self.save_delay_offsets,
1952
        ]:
1953
            try:
×
1954
                method(filename, overwrite)
×
1955
            except (ValueError, AttributeError, KeyError):
×
1956
                pass
×
1957

1958
    def add_jitter(
1✔
1959
        self,
1960
        cols: List[str] = None,
1961
        amps: Union[float, Sequence[float]] = None,
1962
        **kwds,
1963
    ):
1964
        """Add jitter to the selected dataframe columns.
1965

1966
        Args:
1967
            cols (List[str], optional): The columns onto which to apply jitter.
1968
                Defaults to config["dataframe"]["jitter_cols"].
1969
            amps (Union[float, Sequence[float]], optional): Amplitude scalings for the
1970
                jittering noise. If one number is given, the same is used for all axes.
1971
                For uniform noise (default) it will cover the interval [-amp, +amp].
1972
                Defaults to config["dataframe"]["jitter_amps"].
1973
            **kwds: additional keyword arguments passed to ``apply_jitter``.
1974
        """
1975
        if cols is None:
1✔
1976
            cols = self._config["dataframe"]["jitter_cols"]
1✔
1977
        for loc, col in enumerate(cols):
1✔
1978
            if col.startswith("@"):
1✔
1979
                cols[loc] = self._config["dataframe"].get(col.strip("@"))
1✔
1980

1981
        if amps is None:
1✔
1982
            amps = self._config["dataframe"]["jitter_amps"]
1✔
1983

1984
        self._dataframe = self._dataframe.map_partitions(
1✔
1985
            apply_jitter,
1986
            cols=cols,
1987
            cols_jittered=cols,
1988
            amps=amps,
1989
            **kwds,
1990
        )
1991
        if self._timed_dataframe is not None:
1✔
1992
            cols_timed = cols.copy()
1✔
1993
            for col in cols:
1✔
1994
                if col not in self._timed_dataframe.columns:
1✔
1995
                    cols_timed.remove(col)
×
1996

1997
            if cols_timed:
1✔
1998
                self._timed_dataframe = self._timed_dataframe.map_partitions(
1✔
1999
                    apply_jitter,
2000
                    cols=cols_timed,
2001
                    cols_jittered=cols_timed,
2002
                )
2003
        metadata = []
1✔
2004
        for col in cols:
1✔
2005
            metadata.append(col)
1✔
2006
        # TODO: allow only appending if columns are not jittered yet
2007
        self._attributes.add(metadata, "jittering", duplicate_policy="append")
1✔
2008

2009
    def add_time_stamped_data(
1✔
2010
        self,
2011
        dest_column: str,
2012
        time_stamps: np.ndarray = None,
2013
        data: np.ndarray = None,
2014
        archiver_channel: str = None,
2015
        **kwds,
2016
    ):
2017
        """Add data in form of timestamp/value pairs to the dataframe using interpolation to the
2018
        timestamps in the dataframe. The time-stamped data can either be provided, or fetched from
2019
        an EPICS archiver instance.
2020

2021
        Args:
2022
            dest_column (str): destination column name
2023
            time_stamps (np.ndarray, optional): Time stamps of the values to add. If omitted,
2024
                time stamps are retrieved from the epics archiver
2025
            data (np.ndarray, optional): Values corresponding at the time stamps in time_stamps.
2026
                If omitted, data are retrieved from the epics archiver.
2027
            archiver_channel (str, optional): EPICS archiver channel from which to retrieve data.
2028
                Either this or data and time_stamps have to be present.
2029
            **kwds: additional keyword arguments passed to ``add_time_stamped_data``.
2030
        """
2031
        time_stamp_column = kwds.pop(
1✔
2032
            "time_stamp_column",
2033
            self._config["dataframe"].get("time_stamp_alias", ""),
2034
        )
2035

2036
        if time_stamps is None and data is None:
1✔
2037
            if archiver_channel is None:
×
2038
                raise ValueError(
×
2039
                    "Either archiver_channel or both time_stamps and data have to be present!",
2040
                )
2041
            if self.loader.__name__ != "mpes":
×
2042
                raise NotImplementedError(
×
2043
                    "This function is currently only implemented for the mpes loader!",
2044
                )
2045
            ts_from, ts_to = cast(MpesLoader, self.loader).get_start_and_end_time()
×
2046
            # get channel data with +-5 seconds safety margin
2047
            time_stamps, data = get_archiver_data(
×
2048
                archiver_url=self._config["metadata"].get("archiver_url", ""),
2049
                archiver_channel=archiver_channel,
2050
                ts_from=ts_from - 5,
2051
                ts_to=ts_to + 5,
2052
            )
2053

2054
        self._dataframe = add_time_stamped_data(
1✔
2055
            self._dataframe,
2056
            time_stamps=time_stamps,
2057
            data=data,
2058
            dest_column=dest_column,
2059
            time_stamp_column=time_stamp_column,
2060
            **kwds,
2061
        )
2062
        if self._timed_dataframe is not None:
1✔
2063
            if time_stamp_column in self._timed_dataframe:
1✔
2064
                self._timed_dataframe = add_time_stamped_data(
1✔
2065
                    self._timed_dataframe,
2066
                    time_stamps=time_stamps,
2067
                    data=data,
2068
                    dest_column=dest_column,
2069
                    time_stamp_column=time_stamp_column,
2070
                    **kwds,
2071
                )
2072
        metadata: List[Any] = []
1✔
2073
        metadata.append(dest_column)
1✔
2074
        metadata.append(time_stamps)
1✔
2075
        metadata.append(data)
1✔
2076
        self._attributes.add(metadata, "time_stamped_data", duplicate_policy="append")
1✔
2077

2078
    def pre_binning(
1✔
2079
        self,
2080
        df_partitions: Union[int, Sequence[int]] = 100,
2081
        axes: List[str] = None,
2082
        bins: List[int] = None,
2083
        ranges: Sequence[Tuple[float, float]] = None,
2084
        **kwds,
2085
    ) -> xr.DataArray:
2086
        """Function to do an initial binning of the dataframe loaded to the class.
2087

2088
        Args:
2089
            df_partitions (Union[int, Sequence[int]], optional): Number of dataframe partitions to
2090
                use for the initial binning. Defaults to 100.
2091
            axes (List[str], optional): Axes to bin.
2092
                Defaults to config["momentum"]["axes"].
2093
            bins (List[int], optional): Bin numbers to use for binning.
2094
                Defaults to config["momentum"]["bins"].
2095
            ranges (List[Tuple], optional): Ranges to use for binning.
2096
                Defaults to config["momentum"]["ranges"].
2097
            **kwds: Keyword argument passed to ``compute``.
2098

2099
        Returns:
2100
            xr.DataArray: pre-binned data-array.
2101
        """
2102
        if axes is None:
1✔
2103
            axes = self._config["momentum"]["axes"]
1✔
2104
        for loc, axis in enumerate(axes):
1✔
2105
            if axis.startswith("@"):
1✔
2106
                axes[loc] = self._config["dataframe"].get(axis.strip("@"))
1✔
2107

2108
        if bins is None:
1✔
2109
            bins = self._config["momentum"]["bins"]
1✔
2110
        if ranges is None:
1✔
2111
            ranges_ = list(self._config["momentum"]["ranges"])
1✔
2112
            ranges_[2] = np.asarray(ranges_[2]) / 2 ** (
1✔
2113
                self._config["dataframe"]["tof_binning"] - 1
2114
            )
2115
            ranges = [cast(Tuple[float, float], tuple(v)) for v in ranges_]
1✔
2116

2117
        assert self._dataframe is not None, "dataframe needs to be loaded first!"
1✔
2118

2119
        return self.compute(
1✔
2120
            bins=bins,
2121
            axes=axes,
2122
            ranges=ranges,
2123
            df_partitions=df_partitions,
2124
            **kwds,
2125
        )
2126

2127
    def compute(
1✔
2128
        self,
2129
        bins: Union[
2130
            int,
2131
            dict,
2132
            tuple,
2133
            List[int],
2134
            List[np.ndarray],
2135
            List[tuple],
2136
        ] = 100,
2137
        axes: Union[str, Sequence[str]] = None,
2138
        ranges: Sequence[Tuple[float, float]] = None,
2139
        normalize_to_acquisition_time: Union[bool, str] = False,
2140
        **kwds,
2141
    ) -> xr.DataArray:
2142
        """Compute the histogram along the given dimensions.
2143

2144
        Args:
2145
            bins (int, dict, tuple, List[int], List[np.ndarray], List[tuple], optional):
2146
                Definition of the bins. Can be any of the following cases:
2147

2148
                - an integer describing the number of bins in on all dimensions
2149
                - a tuple of 3 numbers describing start, end and step of the binning
2150
                  range
2151
                - a np.arrays defining the binning edges
2152
                - a list (NOT a tuple) of any of the above (int, tuple or np.ndarray)
2153
                - a dictionary made of the axes as keys and any of the above as values.
2154

2155
                This takes priority over the axes and range arguments. Defaults to 100.
2156
            axes (Union[str, Sequence[str]], optional): The names of the axes (columns)
2157
                on which to calculate the histogram. The order will be the order of the
2158
                dimensions in the resulting array. Defaults to None.
2159
            ranges (Sequence[Tuple[float, float]], optional): list of tuples containing
2160
                the start and end point of the binning range. Defaults to None.
2161
            normalize_to_acquisition_time (Union[bool, str]): Option to normalize the
2162
                result to the acquisition time. If a "slow" axis was scanned, providing
2163
                the name of the scanned axis will compute and apply the corresponding
2164
                normalization histogram. Defaults to False.
2165
            **kwds: Keyword arguments:
2166

2167
                - **hist_mode**: Histogram calculation method. "numpy" or "numba". See
2168
                  ``bin_dataframe`` for details. Defaults to
2169
                  config["binning"]["hist_mode"].
2170
                - **mode**: Defines how the results from each partition are combined.
2171
                  "fast", "lean" or "legacy". See ``bin_dataframe`` for details.
2172
                  Defaults to config["binning"]["mode"].
2173
                - **pbar**: Option to show the tqdm progress bar. Defaults to
2174
                  config["binning"]["pbar"].
2175
                - **n_cores**: Number of CPU cores to use for parallelization.
2176
                  Defaults to config["binning"]["num_cores"] or N_CPU-1.
2177
                - **threads_per_worker**: Limit the number of threads that
2178
                  multiprocessing can spawn per binning thread. Defaults to
2179
                  config["binning"]["threads_per_worker"].
2180
                - **threadpool_api**: The API to use for multiprocessing. "blas",
2181
                  "openmp" or None. See ``threadpool_limit`` for details. Defaults to
2182
                  config["binning"]["threadpool_API"].
2183
                - **df_partitions**: A sequence of dataframe partitions, or the
2184
                  number of the dataframe partitions to use. Defaults to all partitions.
2185
                - **filter**: A Sequence of Dictionaries with entries "col", "lower_bound",
2186
                  "upper_bound" to apply as filter to the dataframe before binning. The
2187
                  dataframe in the class remains unmodified by this.
2188

2189
                Additional kwds are passed to ``bin_dataframe``.
2190

2191
        Raises:
2192
            AssertError: Rises when no dataframe has been loaded.
2193

2194
        Returns:
2195
            xr.DataArray: The result of the n-dimensional binning represented in an
2196
            xarray object, combining the data with the axes.
2197
        """
2198
        assert self._dataframe is not None, "dataframe needs to be loaded first!"
1✔
2199

2200
        hist_mode = kwds.pop("hist_mode", self._config["binning"]["hist_mode"])
1✔
2201
        mode = kwds.pop("mode", self._config["binning"]["mode"])
1✔
2202
        pbar = kwds.pop("pbar", self._config["binning"]["pbar"])
1✔
2203
        num_cores = kwds.pop("num_cores", self._config["binning"]["num_cores"])
1✔
2204
        threads_per_worker = kwds.pop(
1✔
2205
            "threads_per_worker",
2206
            self._config["binning"]["threads_per_worker"],
2207
        )
2208
        threadpool_api = kwds.pop(
1✔
2209
            "threadpool_API",
2210
            self._config["binning"]["threadpool_API"],
2211
        )
2212
        df_partitions: Union[int, Sequence[int]] = kwds.pop("df_partitions", None)
1✔
2213
        if isinstance(df_partitions, int):
1✔
2214
            df_partitions = list(range(0, min(df_partitions, self._dataframe.npartitions)))
1✔
2215
        if df_partitions is not None:
1✔
2216
            dataframe = self._dataframe.partitions[df_partitions]
1✔
2217
        else:
2218
            dataframe = self._dataframe
1✔
2219

2220
        filter_params = kwds.pop("filter", None)
1✔
2221
        if filter_params is not None:
1✔
2222
            try:
1✔
2223
                for param in filter_params:
1✔
2224
                    if "col" not in param:
1✔
2225
                        raise ValueError(
1✔
2226
                            "'col' needs to be defined for each filter entry! ",
2227
                            f"Not present in {param}.",
2228
                        )
2229
                    assert set(param.keys()).issubset({"col", "lower_bound", "upper_bound"})
1✔
2230
                    dataframe = apply_filter(dataframe, **param)
1✔
2231
            except AssertionError as exc:
1✔
2232
                invalid_keys = set(param.keys()) - {"lower_bound", "upper_bound"}
1✔
2233
                raise ValueError(
1✔
2234
                    "Only 'col', 'lower_bound' and 'upper_bound' allowed as filter entries. ",
2235
                    f"Parameters {invalid_keys} not valid in {param}.",
2236
                ) from exc
2237

2238
        self._binned = bin_dataframe(
1✔
2239
            df=dataframe,
2240
            bins=bins,
2241
            axes=axes,
2242
            ranges=ranges,
2243
            hist_mode=hist_mode,
2244
            mode=mode,
2245
            pbar=pbar,
2246
            n_cores=num_cores,
2247
            threads_per_worker=threads_per_worker,
2248
            threadpool_api=threadpool_api,
2249
            **kwds,
2250
        )
2251

2252
        for dim in self._binned.dims:
1✔
2253
            try:
1✔
2254
                self._binned[dim].attrs["unit"] = self._config["dataframe"]["units"][dim]
1✔
2255
            except KeyError:
1✔
2256
                pass
1✔
2257

2258
        self._binned.attrs["units"] = "counts"
1✔
2259
        self._binned.attrs["long_name"] = "photoelectron counts"
1✔
2260
        self._binned.attrs["metadata"] = self._attributes.metadata
1✔
2261

2262
        if normalize_to_acquisition_time:
1✔
2263
            if isinstance(normalize_to_acquisition_time, str):
1✔
2264
                axis = normalize_to_acquisition_time
1✔
2265
                print(
1✔
2266
                    f"Calculate normalization histogram for axis '{axis}'...",
2267
                )
2268
                self._normalization_histogram = self.get_normalization_histogram(
1✔
2269
                    axis=axis,
2270
                    df_partitions=df_partitions,
2271
                )
2272
                # if the axes are named correctly, xarray figures out the normalization correctly
2273
                self._normalized = self._binned / self._normalization_histogram
1✔
2274
                self._attributes.add(
1✔
2275
                    self._normalization_histogram.values,
2276
                    name="normalization_histogram",
2277
                    duplicate_policy="overwrite",
2278
                )
2279
            else:
2280
                acquisition_time = self.loader.get_elapsed_time(
×
2281
                    fids=df_partitions,
2282
                )
2283
                if acquisition_time > 0:
×
2284
                    self._normalized = self._binned / acquisition_time
×
2285
                self._attributes.add(
×
2286
                    acquisition_time,
2287
                    name="normalization_histogram",
2288
                    duplicate_policy="overwrite",
2289
                )
2290

2291
            self._normalized.attrs["units"] = "counts/second"
1✔
2292
            self._normalized.attrs["long_name"] = "photoelectron counts per second"
1✔
2293
            self._normalized.attrs["metadata"] = self._attributes.metadata
1✔
2294

2295
            return self._normalized
1✔
2296

2297
        return self._binned
1✔
2298

2299
    def get_normalization_histogram(
1✔
2300
        self,
2301
        axis: str = "delay",
2302
        use_time_stamps: bool = False,
2303
        **kwds,
2304
    ) -> xr.DataArray:
2305
        """Generates a normalization histogram from the timed dataframe. Optionally,
2306
        use the TimeStamps column instead.
2307

2308
        Args:
2309
            axis (str, optional): The axis for which to compute histogram.
2310
                Defaults to "delay".
2311
            use_time_stamps (bool, optional): Use the TimeStamps column of the
2312
                dataframe, rather than the timed dataframe. Defaults to False.
2313
            **kwds: Keyword arguments:
2314

2315
                - **df_partitions**: A sequence of dataframe partitions, or the
2316
                  number of the dataframe partitions to use. Defaults to all partitions.
2317

2318
        Raises:
2319
            ValueError: Raised if no data are binned.
2320
            ValueError: Raised if 'axis' not in binned coordinates.
2321
            ValueError: Raised if config["dataframe"]["time_stamp_alias"] not found
2322
                in Dataframe.
2323

2324
        Returns:
2325
            xr.DataArray: The computed normalization histogram (in TimeStamp units
2326
            per bin).
2327
        """
2328

2329
        if self._binned is None:
1✔
2330
            raise ValueError("Need to bin data first!")
1✔
2331
        if axis not in self._binned.coords:
1✔
2332
            raise ValueError(f"Axis '{axis}' not found in binned data!")
1✔
2333

2334
        df_partitions: Union[int, Sequence[int]] = kwds.pop("df_partitions", None)
1✔
2335
        if isinstance(df_partitions, int):
1✔
2336
            df_partitions = list(range(0, min(df_partitions, self._dataframe.npartitions)))
1✔
2337
        if use_time_stamps or self._timed_dataframe is None:
1✔
2338
            if df_partitions is not None:
1✔
2339
                self._normalization_histogram = normalization_histogram_from_timestamps(
1✔
2340
                    self._dataframe.partitions[df_partitions],
2341
                    axis,
2342
                    self._binned.coords[axis].values,
2343
                    self._config["dataframe"]["time_stamp_alias"],
2344
                )
2345
            else:
2346
                self._normalization_histogram = normalization_histogram_from_timestamps(
×
2347
                    self._dataframe,
2348
                    axis,
2349
                    self._binned.coords[axis].values,
2350
                    self._config["dataframe"]["time_stamp_alias"],
2351
                )
2352
        else:
2353
            if df_partitions is not None:
1✔
2354
                self._normalization_histogram = normalization_histogram_from_timed_dataframe(
1✔
2355
                    self._timed_dataframe.partitions[df_partitions],
2356
                    axis,
2357
                    self._binned.coords[axis].values,
2358
                    self._config["dataframe"]["timed_dataframe_unit_time"],
2359
                )
2360
            else:
2361
                self._normalization_histogram = normalization_histogram_from_timed_dataframe(
×
2362
                    self._timed_dataframe,
2363
                    axis,
2364
                    self._binned.coords[axis].values,
2365
                    self._config["dataframe"]["timed_dataframe_unit_time"],
2366
                )
2367

2368
        return self._normalization_histogram
1✔
2369

2370
    def view_event_histogram(
1✔
2371
        self,
2372
        dfpid: int,
2373
        ncol: int = 2,
2374
        bins: Sequence[int] = None,
2375
        axes: Sequence[str] = None,
2376
        ranges: Sequence[Tuple[float, float]] = None,
2377
        backend: str = "bokeh",
2378
        legend: bool = True,
2379
        histkwds: dict = None,
2380
        legkwds: dict = None,
2381
        **kwds,
2382
    ):
2383
        """Plot individual histograms of specified dimensions (axes) from a substituent
2384
        dataframe partition.
2385

2386
        Args:
2387
            dfpid (int): Number of the data frame partition to look at.
2388
            ncol (int, optional): Number of columns in the plot grid. Defaults to 2.
2389
            bins (Sequence[int], optional): Number of bins to use for the specified
2390
                axes. Defaults to config["histogram"]["bins"].
2391
            axes (Sequence[str], optional): Names of the axes to display.
2392
                Defaults to config["histogram"]["axes"].
2393
            ranges (Sequence[Tuple[float, float]], optional): Value ranges of all
2394
                specified axes. Defaults to config["histogram"]["ranges"].
2395
            backend (str, optional): Backend of the plotting library
2396
                ('matplotlib' or 'bokeh'). Defaults to "bokeh".
2397
            legend (bool, optional): Option to include a legend in the histogram plots.
2398
                Defaults to True.
2399
            histkwds (dict, optional): Keyword arguments for histograms
2400
                (see ``matplotlib.pyplot.hist()``). Defaults to {}.
2401
            legkwds (dict, optional): Keyword arguments for legend
2402
                (see ``matplotlib.pyplot.legend()``). Defaults to {}.
2403
            **kwds: Extra keyword arguments passed to
2404
                ``sed.diagnostics.grid_histogram()``.
2405

2406
        Raises:
2407
            TypeError: Raises when the input values are not of the correct type.
2408
        """
2409
        if bins is None:
1✔
2410
            bins = self._config["histogram"]["bins"]
1✔
2411
        if axes is None:
1✔
2412
            axes = self._config["histogram"]["axes"]
1✔
2413
        axes = list(axes)
1✔
2414
        for loc, axis in enumerate(axes):
1✔
2415
            if axis.startswith("@"):
1✔
2416
                axes[loc] = self._config["dataframe"].get(axis.strip("@"))
1✔
2417
        if ranges is None:
1✔
2418
            ranges = list(self._config["histogram"]["ranges"])
1✔
2419
            for loc, axis in enumerate(axes):
1✔
2420
                if axis == self._config["dataframe"]["tof_column"]:
1✔
2421
                    ranges[loc] = np.asarray(ranges[loc]) / 2 ** (
1✔
2422
                        self._config["dataframe"]["tof_binning"] - 1
2423
                    )
2424
                elif axis == self._config["dataframe"]["adc_column"]:
1✔
2425
                    ranges[loc] = np.asarray(ranges[loc]) / 2 ** (
×
2426
                        self._config["dataframe"]["adc_binning"] - 1
2427
                    )
2428

2429
        input_types = map(type, [axes, bins, ranges])
1✔
2430
        allowed_types = [list, tuple]
1✔
2431

2432
        df = self._dataframe
1✔
2433

2434
        if not set(input_types).issubset(allowed_types):
1✔
2435
            raise TypeError(
×
2436
                "Inputs of axes, bins, ranges need to be list or tuple!",
2437
            )
2438

2439
        # Read out the values for the specified groups
2440
        group_dict_dd = {}
1✔
2441
        dfpart = df.get_partition(dfpid)
1✔
2442
        cols = dfpart.columns
1✔
2443
        for ax in axes:
1✔
2444
            group_dict_dd[ax] = dfpart.values[:, cols.get_loc(ax)]
1✔
2445
        group_dict = ddf.compute(group_dict_dd)[0]
1✔
2446

2447
        # Plot multiple histograms in a grid
2448
        grid_histogram(
1✔
2449
            group_dict,
2450
            ncol=ncol,
2451
            rvs=axes,
2452
            rvbins=bins,
2453
            rvranges=ranges,
2454
            backend=backend,
2455
            legend=legend,
2456
            histkwds=histkwds,
2457
            legkwds=legkwds,
2458
            **kwds,
2459
        )
2460

2461
    def save(
1✔
2462
        self,
2463
        faddr: str,
2464
        **kwds,
2465
    ):
2466
        """Saves the binned data to the provided path and filename.
2467

2468
        Args:
2469
            faddr (str): Path and name of the file to write. Its extension determines
2470
                the file type to write. Valid file types are:
2471

2472
                - "*.tiff", "*.tif": Saves a TIFF stack.
2473
                - "*.h5", "*.hdf5": Saves an HDF5 file.
2474
                - "*.nxs", "*.nexus": Saves a NeXus file.
2475

2476
            **kwds: Keyword arguments, which are passed to the writer functions:
2477
                For TIFF writing:
2478

2479
                - **alias_dict**: Dictionary of dimension aliases to use.
2480

2481
                For HDF5 writing:
2482

2483
                - **mode**: hdf5 read/write mode. Defaults to "w".
2484

2485
                For NeXus:
2486

2487
                - **reader**: Name of the pynxtools reader to use.
2488
                  Defaults to config["nexus"]["reader"]
2489
                - **definition**: NeXus application definition to use for saving.
2490
                  Must be supported by the used ``reader``. Defaults to
2491
                  config["nexus"]["definition"]
2492
                - **input_files**: A list of input files to pass to the reader.
2493
                  Defaults to config["nexus"]["input_files"]
2494
                - **eln_data**: An electronic-lab-notebook file in '.yaml' format
2495
                  to add to the list of files to pass to the reader.
2496
        """
2497
        if self._binned is None:
1✔
2498
            raise NameError("Need to bin data first!")
1✔
2499

2500
        if self._normalized is not None:
1✔
2501
            data = self._normalized
×
2502
        else:
2503
            data = self._binned
1✔
2504

2505
        extension = pathlib.Path(faddr).suffix
1✔
2506

2507
        if extension in (".tif", ".tiff"):
1✔
2508
            to_tiff(
1✔
2509
                data=data,
2510
                faddr=faddr,
2511
                **kwds,
2512
            )
2513
        elif extension in (".h5", ".hdf5"):
1✔
2514
            to_h5(
1✔
2515
                data=data,
2516
                faddr=faddr,
2517
                **kwds,
2518
            )
2519
        elif extension in (".nxs", ".nexus"):
1✔
2520
            try:
1✔
2521
                reader = kwds.pop("reader", self._config["nexus"]["reader"])
1✔
2522
                definition = kwds.pop(
1✔
2523
                    "definition",
2524
                    self._config["nexus"]["definition"],
2525
                )
2526
                input_files = kwds.pop(
1✔
2527
                    "input_files",
2528
                    self._config["nexus"]["input_files"],
2529
                )
2530
            except KeyError as exc:
×
2531
                raise ValueError(
×
2532
                    "The nexus reader, definition and input files need to be provide!",
2533
                ) from exc
2534

2535
            if isinstance(input_files, str):
1✔
2536
                input_files = [input_files]
1✔
2537

2538
            if "eln_data" in kwds:
1✔
2539
                input_files.append(kwds.pop("eln_data"))
1✔
2540

2541
            to_nexus(
1✔
2542
                data=data,
2543
                faddr=faddr,
2544
                reader=reader,
2545
                definition=definition,
2546
                input_files=input_files,
2547
                **kwds,
2548
            )
2549

2550
        else:
2551
            raise NotImplementedError(
1✔
2552
                f"Unrecognized file format: {extension}.",
2553
            )
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc