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

OpenCOMPES / sed / 8957493720

05 May 2024 10:16AM UTC coverage: 91.251% (-0.3%) from 91.532%
8957493720

Pull #395

github

web-flow
Merge c65194a30 into c897413c1
Pull Request #395: Html representation of processor and metadata in notebooks

16 of 41 new or added lines in 3 files covered. (39.02%)

17 existing lines in 1 file now uncovered.

6049 of 6629 relevant lines covered (91.25%)

0.91 hits per line

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

86.43
/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 = "Data Frame: 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✔
NEW
185
        html = "<div>"
×
186

NEW
187
        html += (
×
188
            f"<details><summary>Dataframe</summary>{self.dataframe.head()._repr_html_()}</details>"
189
        )
190

191
        # Add expandable section for dataframe
NEW
192
        html += f"<details><summary>Dask</summary>{self.dataframe._repr_html_()}</details>"
×
193

194
        # Add expandable section for attributes
NEW
195
        html += "<details><summary>Metadata</summary>"
×
NEW
196
        html += self.attributes._repr_html_()
×
NEW
197
        html += "</details>"
×
198

199
        # Add expandable section for plots
NEW
200
        html += "<details><summary>Plots</summary>"
×
201
        # Something like the event histogram can be added here,
202
        # but the method needs to output image/html
203
        # self.view_event_histogram(dfpid=2, backend="matplotlib")
NEW
204
        html += "</details>"
×
205

NEW
206
        html += "</div>"
×
207

NEW
208
        return html
×
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 normalizazion 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
        ptovided 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 algortihm.
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"]["transformtions"].
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 homomorphy 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.slice_corrected = self.mc.slice
1✔
715

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1445
        tof_column = self._config["dataframe"]["tof_column"]
1✔
1446

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1606
        """
1607
        if verbose is None:
1✔
1608
            verbose = self.verbose
1✔
1609

1610
        tof_column = self._config["dataframe"]["tof_column"]
1✔
1611

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

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

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

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

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

1664
        tof_column = self._config["dataframe"]["tof_column"]
1✔
1665

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

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

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

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

1711
        Args:
1712
            delay_range (Tuple[float, float], optional): The scanned delay range in
1713
                picoseconds. Defaults to None.
1714
            datafile (str, optional): The file from which to read the delay ranges.
1715
                Defaults to None.
1716
            preview (bool, optional): Option to preview the first elements of the data frame.
1717
                Defaults to False.
1718
            verbose (bool, optional): Option to print out diagnostic information.
1719
                Defaults to config["core"]["verbose"].
1720
            **kwds: Keyword args passed to ``DelayCalibrator.append_delay_axis``.
1721
        """
1722
        if verbose is None:
1✔
1723
            verbose = self.verbose
1✔
1724

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

1729
        if self._dataframe is not None:
1✔
1730
            if verbose:
1✔
1731
                print("Adding delay column to dataframe:")
1✔
1732

1733
            if delay_range is None and datafile is None:
1✔
1734
                if len(self.dc.calibration) == 0:
1✔
1735
                    try:
1✔
1736
                        datafile = self._files[0]
1✔
1737
                    except IndexError:
×
1738
                        print(
×
1739
                            "No datafile available, specify either",
1740
                            " 'datafile' or 'delay_range'",
1741
                        )
1742
                        raise
×
1743

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

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

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

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

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

1804
        if "creation_date" not in calibration:
1✔
1805
            calibration["creation_date"] = datetime.now().timestamp()
×
1806

1807
        config = {
1✔
1808
            "delay": {
1809
                "calibration": calibration,
1810
            },
1811
        }
1812
        save_config(config, filename, overwrite)
1✔
1813

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

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

1844
        Raises:
1845
            ValueError: If the delay column is not in the dataframe.
1846
        """
1847
        if verbose is None:
1✔
1848
            verbose = self.verbose
1✔
1849

1850
        delay_column = self._config["dataframe"]["delay_column"]
1✔
1851
        if delay_column not in self._dataframe.columns:
1✔
1852
            raise ValueError(f"Delay column {delay_column} not found in dataframe! ")
1✔
1853

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

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

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

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

1915
        if "creation_date" not in self.ec.offsets.keys():
1✔
1916
            self.ec.offsets["creation_date"] = datetime.now().timestamp()
1✔
1917

1918
        config = {
1✔
1919
            "delay": {
1920
                "offsets": self.dc.offsets,
1921
            },
1922
        }
1923
        save_config(config, filename, overwrite)
1✔
1924
        print(f'Saved delay offset parameters to "{filename}".')
1✔
1925

1926
    def save_workflow_params(
1✔
1927
        self,
1928
        filename: str = None,
1929
        overwrite: bool = False,
1930
    ) -> None:
1931
        """run all save calibration parameter methods
1932

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

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

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

1977
        if amps is None:
1✔
1978
            amps = self._config["dataframe"]["jitter_amps"]
1✔
1979

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

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

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

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

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

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

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

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

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

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

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

2115
        return self.compute(
1✔
2116
            bins=bins,
2117
            axes=axes,
2118
            ranges=ranges,
2119
            df_partitions=df_partitions,
2120
            **kwds,
2121
        )
2122

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

2140
        Args:
2141
            bins (int, dict, tuple, List[int], List[np.ndarray], List[tuple], optional):
2142
                Definition of the bins. Can be any of the following cases:
2143

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

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

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

2185
                Additional kwds are passed to ``bin_dataframe``.
2186

2187
        Raises:
2188
            AssertError: Rises when no dataframe has been loaded.
2189

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

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

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

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

2248
        for dim in self._binned.dims:
1✔
2249
            try:
1✔
2250
                self._binned[dim].attrs["unit"] = self._config["dataframe"]["units"][dim]
1✔
2251
            except KeyError:
1✔
2252
                pass
1✔
2253

2254
        self._binned.attrs["units"] = "counts"
1✔
2255
        self._binned.attrs["long_name"] = "photoelectron counts"
1✔
2256
        self._binned.attrs["metadata"] = self._attributes.metadata
1✔
2257

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

2287
            self._normalized.attrs["units"] = "counts/second"
1✔
2288
            self._normalized.attrs["long_name"] = "photoelectron counts per second"
1✔
2289
            self._normalized.attrs["metadata"] = self._attributes.metadata
1✔
2290

2291
            return self._normalized
1✔
2292

2293
        return self._binned
1✔
2294

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

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

2311
                - **df_partitions**: A sequence of dataframe partitions, or the
2312
                  number of the dataframe partitions to use. Defaults to all partitions.
2313

2314
        Raises:
2315
            ValueError: Raised if no data are binned.
2316
            ValueError: Raised if 'axis' not in binned coordinates.
2317
            ValueError: Raised if config["dataframe"]["time_stamp_alias"] not found
2318
                in Dataframe.
2319

2320
        Returns:
2321
            xr.DataArray: The computed normalization histogram (in TimeStamp units
2322
            per bin).
2323
        """
2324

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

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

2364
        return self._normalization_histogram
1✔
2365

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

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

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

2425
        input_types = map(type, [axes, bins, ranges])
1✔
2426
        allowed_types = [list, tuple]
1✔
2427

2428
        df = self._dataframe
1✔
2429

2430
        if not set(input_types).issubset(allowed_types):
1✔
2431
            raise TypeError(
×
2432
                "Inputs of axes, bins, ranges need to be list or tuple!",
2433
            )
2434

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

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

2457
    def save(
1✔
2458
        self,
2459
        faddr: str,
2460
        **kwds,
2461
    ):
2462
        """Saves the binned data to the provided path and filename.
2463

2464
        Args:
2465
            faddr (str): Path and name of the file to write. Its extension determines
2466
                the file type to write. Valid file types are:
2467

2468
                - "*.tiff", "*.tif": Saves a TIFF stack.
2469
                - "*.h5", "*.hdf5": Saves an HDF5 file.
2470
                - "*.nxs", "*.nexus": Saves a NeXus file.
2471

2472
            **kwds: Keyword argumens, which are passed to the writer functions:
2473
                For TIFF writing:
2474

2475
                - **alias_dict**: Dictionary of dimension aliases to use.
2476

2477
                For HDF5 writing:
2478

2479
                - **mode**: hdf5 read/write mode. Defaults to "w".
2480

2481
                For NeXus:
2482

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

2496
        if self._normalized is not None:
1✔
2497
            data = self._normalized
×
2498
        else:
2499
            data = self._binned
1✔
2500

2501
        extension = pathlib.Path(faddr).suffix
1✔
2502

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

2531
            if isinstance(input_files, str):
1✔
2532
                input_files = [input_files]
1✔
2533

2534
            if "eln_data" in kwds:
1✔
2535
                input_files.append(kwds.pop("eln_data"))
×
2536

2537
            to_nexus(
1✔
2538
                data=data,
2539
                faddr=faddr,
2540
                reader=reader,
2541
                definition=definition,
2542
                input_files=input_files,
2543
                **kwds,
2544
            )
2545

2546
        else:
2547
            raise NotImplementedError(
1✔
2548
                f"Unrecognized file format: {extension}.",
2549
            )
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