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

OpenCOMPES / sed / 9553353548

14 Jun 2024 08:58PM UTC coverage: 91.952% (-0.01%) from 91.962%
9553353548

Pull #411

github

rettigl
Merge remote-tracking branch 'origin/v1_feature_branch' into energy_calibration_bias_shift
Pull Request #411: Energy calibration bias shift

76 of 87 new or added lines in 3 files covered. (87.36%)

3 existing lines in 1 file now uncovered.

6490 of 7058 relevant lines covered (91.95%)

0.92 hits per line

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

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

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

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

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

42
N_CPU = psutil.cpu_count()
1✔
43

44

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

202
        return html
×
203

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

333
        Returns:
334
            xr.DataArray: The 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_energy: float,
1295
        method: str = None,
1296
        energy_scale: str = None,
1297
        verbose: bool = None,
1298
        **kwds,
1299
    ):
1300
        """3. Step of the energy calibration workflow: Calculate the calibration
1301
        function for the energy axis, and apply it to the dataframe. Two
1302
        approximations are implemented, a (normally 3rd order) polynomial
1303
        approximation, and a d^2/(t-t0)^2 relation.
1304

1305
        Args:
1306
            ref_energy (float): Binding/kinetic energy of the detected feature.
1307
            method (str, optional): Method for determining the energy calibration.
1308

1309
                - **'lmfit'**: Energy calibration using lmfit and 1/t^2 form.
1310
                - **'lstsq'**, **'lsqr'**: Energy calibration using polynomial form.
1311

1312
                Defaults to config["energy"]["calibration_method"]
1313
            energy_scale (str, optional): Direction of increasing energy scale.
1314

1315
                - **'kinetic'**: increasing energy with decreasing TOF.
1316
                - **'binding'**: increasing energy with increasing TOF.
1317

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

1326
        if method is None:
1✔
1327
            method = self._config["energy"]["calibration_method"]
1✔
1328

1329
        if energy_scale is None:
1✔
1330
            energy_scale = self._config["energy"]["energy_scale"]
1✔
1331

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

1378
    # 3a. Save energy calibration parameters to config file.
1379
    def save_energy_calibration(
1✔
1380
        self,
1381
        filename: str = None,
1382
        overwrite: bool = False,
1383
    ):
1384
        """Save the generated energy calibration parameters to the folder config file.
1385

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

1407
        if "creation_date" not in calibration:
1✔
1408
            calibration["creation_date"] = datetime.now().timestamp()
×
1409

1410
        config = {"energy": {"calibration": calibration}}
1✔
1411
        save_config(config, filename, overwrite)
1✔
1412
        print(f'Saved energy calibration parameters to "{filename}".')
1✔
1413

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

1428
        Args:
1429
            calibration (dict, optional): Calibration dict containing calibration
1430
                parameters. Overrides calibration from class or config.
1431
                Defaults to None.
1432
            bias_voltage (float, optional): Sample bias voltage of the scan data. If omitted,
1433
                the bias voltage is being read from the dataframe. If it is not found there,
1434
                a warning is printed and the calibrated data will not be offset correctly.
1435
            preview (bool): Option to preview the first elements of the data frame.
1436
            verbose (bool, optional): Option to print out diagnostic information.
1437
                Defaults to config["core"]["verbose"].
1438
            **kwds:
1439
                Keyword args passed to ``EnergyCalibrator.append_energy_axis()``.
1440
        """
1441
        if verbose is None:
1✔
1442
            verbose = self.verbose
1✔
1443

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

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

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

1473
        else:
1474
            raise ValueError("No dataframe loaded!")
×
1475

1476
        if bias_voltage is not None:
1✔
NEW
1477
            self.add_energy_offset(constant=bias_voltage, verbose=verbose, preview=preview)
×
1478
        elif self.config["dataframe"]["bias_column"] in self._dataframe.columns:
1✔
1479
            self.add_energy_offset(
1✔
1480
                columns=[self.config["dataframe"]["bias_column"]],
1481
                verbose=verbose,
1482
                preview=preview,
1483
            )
1484
        else:
1485
            print("Sample bias data not found or provided. Calibrated energy will be offset.")
1✔
1486
            # Preview only if no offset applied
1487
            if preview:
1✔
NEW
1488
                print(self._dataframe.head(10))
×
1489
            else:
1490
                if verbose:
1✔
1491
                    print(self._dataframe)
1✔
1492

1493
    def add_energy_offset(
1✔
1494
        self,
1495
        constant: float = None,
1496
        columns: Union[str, Sequence[str]] = None,
1497
        weights: Union[float, Sequence[float]] = None,
1498
        reductions: Union[str, Sequence[str]] = None,
1499
        preserve_mean: Union[bool, Sequence[bool]] = None,
1500
        preview: bool = False,
1501
        verbose: bool = None,
1502
    ) -> None:
1503
        """Shift the energy axis of the dataframe by a given amount.
1504

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

1521
        Raises:
1522
            ValueError: If the energy column is not in the dataframe.
1523
        """
1524
        if verbose is None:
1✔
1525
            verbose = self.verbose
1✔
1526

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

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

1574
    def save_energy_offset(
1✔
1575
        self,
1576
        filename: str = None,
1577
        overwrite: bool = False,
1578
    ):
1579
        """Save the generated energy calibration parameters to the folder config file.
1580

1581
        Args:
1582
            filename (str, optional): Filename of the config dictionary to save to.
1583
                Defaults to "sed_config.yaml" in the current folder.
1584
            overwrite (bool, optional): Option to overwrite the present dictionary.
1585
                Defaults to False.
1586
        """
1587
        if filename is None:
×
1588
            filename = "sed_config.yaml"
×
1589
        if len(self.ec.offsets) == 0:
×
1590
            raise ValueError("No energy offset parameters to save!")
×
1591

1592
        if "creation_date" not in self.ec.offsets.keys():
×
1593
            self.ec.offsets["creation_date"] = datetime.now().timestamp()
×
1594

1595
        config = {"energy": {"offsets": self.ec.offsets}}
×
1596
        save_config(config, filename, overwrite)
×
1597
        print(f'Saved energy offset parameters to "{filename}".')
×
1598

1599
    def append_tof_ns_axis(
1✔
1600
        self,
1601
        preview: bool = False,
1602
        verbose: bool = None,
1603
        **kwds,
1604
    ):
1605
        """Convert time-of-flight channel steps to nanoseconds.
1606

1607
        Args:
1608
            tof_ns_column (str, optional): Name of the generated column containing the
1609
                time-of-flight in nanosecond.
1610
                Defaults to config["dataframe"]["tof_ns_column"].
1611
            preview (bool, optional): Option to preview the first elements of the data frame.
1612
                Defaults to False.
1613
            verbose (bool, optional): Option to print out diagnostic information.
1614
                Defaults to config["core"]["verbose"].
1615
            **kwds: additional arguments are passed to ``EnergyCalibrator.tof_step_to_ns()``.
1616

1617
        """
1618
        if verbose is None:
1✔
1619
            verbose = self.verbose
1✔
1620

1621
        tof_column = self._config["dataframe"]["tof_column"]
1✔
1622

1623
        if self._dataframe is not None:
1✔
1624
            if verbose:
1✔
1625
                print("Adding time-of-flight column in nanoseconds to dataframe:")
1✔
1626
            # TODO assert order of execution through metadata
1627

1628
            df, metadata = self.ec.append_tof_ns_axis(
1✔
1629
                df=self._dataframe,
1630
                **kwds,
1631
            )
1632
            if self._timed_dataframe is not None and tof_column in self._timed_dataframe.columns:
1✔
1633
                tdf, _ = self.ec.append_tof_ns_axis(
1✔
1634
                    df=self._timed_dataframe,
1635
                    **kwds,
1636
                )
1637

1638
            self._attributes.add(
1✔
1639
                metadata,
1640
                "tof_ns_conversion",
1641
                duplicate_policy="overwrite",
1642
            )
1643
            self._dataframe = df
1✔
1644
            if self._timed_dataframe is not None and tof_column in self._timed_dataframe.columns:
1✔
1645
                self._timed_dataframe = tdf
1✔
1646
        else:
1647
            raise ValueError("No dataframe loaded!")
×
1648
        if preview:
1✔
1649
            print(self._dataframe.head(10))
×
1650
        else:
1651
            if verbose:
1✔
1652
                print(self._dataframe)
1✔
1653

1654
    def align_dld_sectors(
1✔
1655
        self,
1656
        sector_delays: np.ndarray = None,
1657
        preview: bool = False,
1658
        verbose: bool = None,
1659
        **kwds,
1660
    ):
1661
        """Align the 8s sectors of the HEXTOF endstation.
1662

1663
        Args:
1664
            sector_delays (np.ndarray, optional): Array containing the sector delays. Defaults to
1665
                config["dataframe"]["sector_delays"].
1666
            preview (bool, optional): Option to preview the first elements of the data frame.
1667
                Defaults to False.
1668
            verbose (bool, optional): Option to print out diagnostic information.
1669
                Defaults to config["core"]["verbose"].
1670
            **kwds: additional arguments are passed to ``EnergyCalibrator.align_dld_sectors()``.
1671
        """
1672
        if verbose is None:
1✔
1673
            verbose = self.verbose
1✔
1674

1675
        tof_column = self._config["dataframe"]["tof_column"]
1✔
1676

1677
        if self._dataframe is not None:
1✔
1678
            if verbose:
1✔
1679
                print("Aligning 8s sectors of dataframe")
1✔
1680
            # TODO assert order of execution through metadata
1681

1682
            df, metadata = self.ec.align_dld_sectors(
1✔
1683
                df=self._dataframe,
1684
                sector_delays=sector_delays,
1685
                **kwds,
1686
            )
1687
            if self._timed_dataframe is not None and tof_column in self._timed_dataframe.columns:
1✔
1688
                tdf, _ = self.ec.align_dld_sectors(
×
1689
                    df=self._timed_dataframe,
1690
                    sector_delays=sector_delays,
1691
                    **kwds,
1692
                )
1693

1694
            self._attributes.add(
1✔
1695
                metadata,
1696
                "dld_sector_alignment",
1697
                duplicate_policy="raise",
1698
            )
1699
            self._dataframe = df
1✔
1700
            if self._timed_dataframe is not None and tof_column in self._timed_dataframe.columns:
1✔
1701
                self._timed_dataframe = tdf
×
1702
        else:
1703
            raise ValueError("No dataframe loaded!")
×
1704
        if preview:
1✔
1705
            print(self._dataframe.head(10))
×
1706
        else:
1707
            if verbose:
1✔
1708
                print(self._dataframe)
1✔
1709

1710
    # Delay calibration function
1711
    def calibrate_delay_axis(
1✔
1712
        self,
1713
        delay_range: Tuple[float, float] = None,
1714
        datafile: str = None,
1715
        preview: bool = False,
1716
        verbose: bool = None,
1717
        **kwds,
1718
    ):
1719
        """Append delay column to dataframe. Either provide delay ranges, or read
1720
        them from a file.
1721

1722
        Args:
1723
            delay_range (Tuple[float, float], optional): The scanned delay range in
1724
                picoseconds. Defaults to None.
1725
            datafile (str, optional): The file from which to read the delay ranges.
1726
                Defaults to None.
1727
            preview (bool, optional): Option to preview the first elements of the data frame.
1728
                Defaults to False.
1729
            verbose (bool, optional): Option to print out diagnostic information.
1730
                Defaults to config["core"]["verbose"].
1731
            **kwds: Keyword args passed to ``DelayCalibrator.append_delay_axis``.
1732
        """
1733
        if verbose is None:
1✔
1734
            verbose = self.verbose
1✔
1735

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

1740
        if self._dataframe is not None:
1✔
1741
            if verbose:
1✔
1742
                print("Adding delay column to dataframe:")
1✔
1743

1744
            if delay_range is None and datafile is None:
1✔
1745
                if len(self.dc.calibration) == 0:
1✔
1746
                    try:
1✔
1747
                        datafile = self._files[0]
1✔
1748
                    except IndexError:
×
1749
                        print(
×
1750
                            "No datafile available, specify either",
1751
                            " 'datafile' or 'delay_range'",
1752
                        )
1753
                        raise
×
1754

1755
            df, metadata = self.dc.append_delay_axis(
1✔
1756
                self._dataframe,
1757
                delay_range=delay_range,
1758
                datafile=datafile,
1759
                verbose=verbose,
1760
                **kwds,
1761
            )
1762
            if self._timed_dataframe is not None and adc_column in self._timed_dataframe.columns:
1✔
1763
                tdf, _ = self.dc.append_delay_axis(
1✔
1764
                    self._timed_dataframe,
1765
                    delay_range=delay_range,
1766
                    datafile=datafile,
1767
                    verbose=False,
1768
                    **kwds,
1769
                )
1770

1771
            # Add Metadata
1772
            self._attributes.add(
1✔
1773
                metadata,
1774
                "delay_calibration",
1775
                duplicate_policy="overwrite",
1776
            )
1777
            self._dataframe = df
1✔
1778
            if self._timed_dataframe is not None and adc_column in self._timed_dataframe.columns:
1✔
1779
                self._timed_dataframe = tdf
1✔
1780
        else:
1781
            raise ValueError("No dataframe loaded!")
×
1782
        if preview:
1✔
1783
            print(self._dataframe.head(10))
1✔
1784
        else:
1785
            if self.verbose:
1✔
1786
                print(self._dataframe)
1✔
1787

1788
    def save_delay_calibration(
1✔
1789
        self,
1790
        filename: str = None,
1791
        overwrite: bool = False,
1792
    ) -> None:
1793
        """Save the generated delay calibration parameters to the folder config file.
1794

1795
        Args:
1796
            filename (str, optional): Filename of the config dictionary to save to.
1797
                Defaults to "sed_config.yaml" in the current folder.
1798
            overwrite (bool, optional): Option to overwrite the present dictionary.
1799
                Defaults to False.
1800
        """
1801
        if filename is None:
1✔
1802
            filename = "sed_config.yaml"
×
1803

1804
        if len(self.dc.calibration) == 0:
1✔
1805
            raise ValueError("No delay calibration parameters to save!")
×
1806
        calibration = {}
1✔
1807
        for key, value in self.dc.calibration.items():
1✔
1808
            if key == "datafile":
1✔
1809
                calibration[key] = value
1✔
1810
            elif key in ["adc_range", "delay_range", "delay_range_mm"]:
1✔
1811
                calibration[key] = [float(i) for i in value]
1✔
1812
            else:
1813
                calibration[key] = float(value)
1✔
1814

1815
        if "creation_date" not in calibration:
1✔
1816
            calibration["creation_date"] = datetime.now().timestamp()
×
1817

1818
        config = {
1✔
1819
            "delay": {
1820
                "calibration": calibration,
1821
            },
1822
        }
1823
        save_config(config, filename, overwrite)
1✔
1824

1825
    def add_delay_offset(
1✔
1826
        self,
1827
        constant: float = None,
1828
        flip_delay_axis: bool = None,
1829
        columns: Union[str, Sequence[str]] = None,
1830
        weights: Union[float, Sequence[float]] = 1.0,
1831
        reductions: Union[str, Sequence[str]] = None,
1832
        preserve_mean: Union[bool, Sequence[bool]] = False,
1833
        preview: bool = False,
1834
        verbose: bool = None,
1835
    ) -> None:
1836
        """Shift the delay axis of the dataframe by a constant or other columns.
1837

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

1855
        Raises:
1856
            ValueError: If the delay column is not in the dataframe.
1857
        """
1858
        if verbose is None:
1✔
1859
            verbose = self.verbose
1✔
1860

1861
        delay_column = self._config["dataframe"]["delay_column"]
1✔
1862
        if delay_column not in self._dataframe.columns:
1✔
1863
            raise ValueError(f"Delay column {delay_column} not found in dataframe! ")
1✔
1864

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

1892
            self._attributes.add(
1✔
1893
                metadata,
1894
                "delay_offset",
1895
                duplicate_policy="append",
1896
            )
1897
            self._dataframe = df
1✔
1898
            if self._timed_dataframe is not None and delay_column in self._timed_dataframe.columns:
1✔
1899
                self._timed_dataframe = tdf
1✔
1900
        else:
1901
            raise ValueError("No dataframe loaded!")
×
1902
        if preview:
1✔
1903
            print(self._dataframe.head(10))
1✔
1904
        else:
1905
            if verbose:
1✔
1906
                print(self._dataframe)
1✔
1907

1908
    def save_delay_offsets(
1✔
1909
        self,
1910
        filename: str = None,
1911
        overwrite: bool = False,
1912
    ) -> None:
1913
        """Save the generated delay calibration parameters to the folder config file.
1914

1915
        Args:
1916
            filename (str, optional): Filename of the config dictionary to save to.
1917
                Defaults to "sed_config.yaml" in the current folder.
1918
            overwrite (bool, optional): Option to overwrite the present dictionary.
1919
                Defaults to False.
1920
        """
1921
        if filename is None:
1✔
1922
            filename = "sed_config.yaml"
×
1923
        if len(self.dc.offsets) == 0:
1✔
1924
            raise ValueError("No delay offset parameters to save!")
×
1925

1926
        if "creation_date" not in self.ec.offsets.keys():
1✔
1927
            self.ec.offsets["creation_date"] = datetime.now().timestamp()
1✔
1928

1929
        config = {
1✔
1930
            "delay": {
1931
                "offsets": self.dc.offsets,
1932
            },
1933
        }
1934
        save_config(config, filename, overwrite)
1✔
1935
        print(f'Saved delay offset parameters to "{filename}".')
1✔
1936

1937
    def save_workflow_params(
1✔
1938
        self,
1939
        filename: str = None,
1940
        overwrite: bool = False,
1941
    ) -> None:
1942
        """run all save calibration parameter methods
1943

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

1965
    def add_jitter(
1✔
1966
        self,
1967
        cols: List[str] = None,
1968
        amps: Union[float, Sequence[float]] = None,
1969
        **kwds,
1970
    ):
1971
        """Add jitter to the selected dataframe columns.
1972

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

1988
        if amps is None:
1✔
1989
            amps = self._config["dataframe"]["jitter_amps"]
1✔
1990

1991
        self._dataframe = self._dataframe.map_partitions(
1✔
1992
            apply_jitter,
1993
            cols=cols,
1994
            cols_jittered=cols,
1995
            amps=amps,
1996
            **kwds,
1997
        )
1998
        if self._timed_dataframe is not None:
1✔
1999
            cols_timed = cols.copy()
1✔
2000
            for col in cols:
1✔
2001
                if col not in self._timed_dataframe.columns:
1✔
2002
                    cols_timed.remove(col)
×
2003

2004
            if cols_timed:
1✔
2005
                self._timed_dataframe = self._timed_dataframe.map_partitions(
1✔
2006
                    apply_jitter,
2007
                    cols=cols_timed,
2008
                    cols_jittered=cols_timed,
2009
                )
2010
        metadata = []
1✔
2011
        for col in cols:
1✔
2012
            metadata.append(col)
1✔
2013
        # TODO: allow only appending if columns are not jittered yet
2014
        self._attributes.add(metadata, "jittering", duplicate_policy="append")
1✔
2015

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

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

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

2061
        self._dataframe = add_time_stamped_data(
1✔
2062
            self._dataframe,
2063
            time_stamps=time_stamps,
2064
            data=data,
2065
            dest_column=dest_column,
2066
            time_stamp_column=time_stamp_column,
2067
            **kwds,
2068
        )
2069
        if self._timed_dataframe is not None:
1✔
2070
            if time_stamp_column in self._timed_dataframe:
1✔
2071
                self._timed_dataframe = add_time_stamped_data(
1✔
2072
                    self._timed_dataframe,
2073
                    time_stamps=time_stamps,
2074
                    data=data,
2075
                    dest_column=dest_column,
2076
                    time_stamp_column=time_stamp_column,
2077
                    **kwds,
2078
                )
2079
        metadata: List[Any] = []
1✔
2080
        metadata.append(dest_column)
1✔
2081
        metadata.append(time_stamps)
1✔
2082
        metadata.append(data)
1✔
2083
        self._attributes.add(metadata, "time_stamped_data", duplicate_policy="append")
1✔
2084

2085
    def pre_binning(
1✔
2086
        self,
2087
        df_partitions: Union[int, Sequence[int]] = 100,
2088
        axes: List[str] = None,
2089
        bins: List[int] = None,
2090
        ranges: Sequence[Tuple[float, float]] = None,
2091
        **kwds,
2092
    ) -> xr.DataArray:
2093
        """Function to do an initial binning of the dataframe loaded to the class.
2094

2095
        Args:
2096
            df_partitions (Union[int, Sequence[int]], optional): Number of dataframe partitions to
2097
                use for the initial binning. Defaults to 100.
2098
            axes (List[str], optional): Axes to bin.
2099
                Defaults to config["momentum"]["axes"].
2100
            bins (List[int], optional): Bin numbers to use for binning.
2101
                Defaults to config["momentum"]["bins"].
2102
            ranges (List[Tuple], optional): Ranges to use for binning.
2103
                Defaults to config["momentum"]["ranges"].
2104
            **kwds: Keyword argument passed to ``compute``.
2105

2106
        Returns:
2107
            xr.DataArray: pre-binned data-array.
2108
        """
2109
        if axes is None:
1✔
2110
            axes = self._config["momentum"]["axes"]
1✔
2111
        for loc, axis in enumerate(axes):
1✔
2112
            if axis.startswith("@"):
1✔
2113
                axes[loc] = self._config["dataframe"].get(axis.strip("@"))
1✔
2114

2115
        if bins is None:
1✔
2116
            bins = self._config["momentum"]["bins"]
1✔
2117
        if ranges is None:
1✔
2118
            ranges_ = list(self._config["momentum"]["ranges"])
1✔
2119
            ranges_[2] = np.asarray(ranges_[2]) / 2 ** (
1✔
2120
                self._config["dataframe"]["tof_binning"] - 1
2121
            )
2122
            ranges = [cast(Tuple[float, float], tuple(v)) for v in ranges_]
1✔
2123

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

2126
        return self.compute(
1✔
2127
            bins=bins,
2128
            axes=axes,
2129
            ranges=ranges,
2130
            df_partitions=df_partitions,
2131
            **kwds,
2132
        )
2133

2134
    def compute(
1✔
2135
        self,
2136
        bins: Union[
2137
            int,
2138
            dict,
2139
            tuple,
2140
            List[int],
2141
            List[np.ndarray],
2142
            List[tuple],
2143
        ] = 100,
2144
        axes: Union[str, Sequence[str]] = None,
2145
        ranges: Sequence[Tuple[float, float]] = None,
2146
        normalize_to_acquisition_time: Union[bool, str] = False,
2147
        **kwds,
2148
    ) -> xr.DataArray:
2149
        """Compute the histogram along the given dimensions.
2150

2151
        Args:
2152
            bins (int, dict, tuple, List[int], List[np.ndarray], List[tuple], optional):
2153
                Definition of the bins. Can be any of the following cases:
2154

2155
                - an integer describing the number of bins in on all dimensions
2156
                - a tuple of 3 numbers describing start, end and step of the binning
2157
                  range
2158
                - a np.arrays defining the binning edges
2159
                - a list (NOT a tuple) of any of the above (int, tuple or np.ndarray)
2160
                - a dictionary made of the axes as keys and any of the above as values.
2161

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

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

2196
                Additional kwds are passed to ``bin_dataframe``.
2197

2198
        Raises:
2199
            AssertError: Rises when no dataframe has been loaded.
2200

2201
        Returns:
2202
            xr.DataArray: The result of the n-dimensional binning represented in an
2203
            xarray object, combining the data with the axes.
2204
        """
2205
        assert self._dataframe is not None, "dataframe needs to be loaded first!"
1✔
2206

2207
        hist_mode = kwds.pop("hist_mode", self._config["binning"]["hist_mode"])
1✔
2208
        mode = kwds.pop("mode", self._config["binning"]["mode"])
1✔
2209
        pbar = kwds.pop("pbar", self._config["binning"]["pbar"])
1✔
2210
        num_cores = kwds.pop("num_cores", self._config["binning"]["num_cores"])
1✔
2211
        threads_per_worker = kwds.pop(
1✔
2212
            "threads_per_worker",
2213
            self._config["binning"]["threads_per_worker"],
2214
        )
2215
        threadpool_api = kwds.pop(
1✔
2216
            "threadpool_API",
2217
            self._config["binning"]["threadpool_API"],
2218
        )
2219
        df_partitions: Union[int, Sequence[int]] = kwds.pop("df_partitions", None)
1✔
2220
        if isinstance(df_partitions, int):
1✔
2221
            df_partitions = list(range(0, min(df_partitions, self._dataframe.npartitions)))
1✔
2222
        if df_partitions is not None:
1✔
2223
            dataframe = self._dataframe.partitions[df_partitions]
1✔
2224
        else:
2225
            dataframe = self._dataframe
1✔
2226

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

2245
        self._binned = bin_dataframe(
1✔
2246
            df=dataframe,
2247
            bins=bins,
2248
            axes=axes,
2249
            ranges=ranges,
2250
            hist_mode=hist_mode,
2251
            mode=mode,
2252
            pbar=pbar,
2253
            n_cores=num_cores,
2254
            threads_per_worker=threads_per_worker,
2255
            threadpool_api=threadpool_api,
2256
            **kwds,
2257
        )
2258

2259
        for dim in self._binned.dims:
1✔
2260
            try:
1✔
2261
                self._binned[dim].attrs["unit"] = self._config["dataframe"]["units"][dim]
1✔
2262
            except KeyError:
1✔
2263
                pass
1✔
2264

2265
        self._binned.attrs["units"] = "counts"
1✔
2266
        self._binned.attrs["long_name"] = "photoelectron counts"
1✔
2267
        self._binned.attrs["metadata"] = self._attributes.metadata
1✔
2268

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

2298
            self._normalized.attrs["units"] = "counts/second"
1✔
2299
            self._normalized.attrs["long_name"] = "photoelectron counts per second"
1✔
2300
            self._normalized.attrs["metadata"] = self._attributes.metadata
1✔
2301

2302
            return self._normalized
1✔
2303

2304
        return self._binned
1✔
2305

2306
    def get_normalization_histogram(
1✔
2307
        self,
2308
        axis: str = "delay",
2309
        use_time_stamps: bool = False,
2310
        **kwds,
2311
    ) -> xr.DataArray:
2312
        """Generates a normalization histogram from the timed dataframe. Optionally,
2313
        use the TimeStamps column instead.
2314

2315
        Args:
2316
            axis (str, optional): The axis for which to compute histogram.
2317
                Defaults to "delay".
2318
            use_time_stamps (bool, optional): Use the TimeStamps column of the
2319
                dataframe, rather than the timed dataframe. Defaults to False.
2320
            **kwds: Keyword arguments:
2321

2322
                - **df_partitions**: A sequence of dataframe partitions, or the
2323
                  number of the dataframe partitions to use. Defaults to all partitions.
2324

2325
        Raises:
2326
            ValueError: Raised if no data are binned.
2327
            ValueError: Raised if 'axis' not in binned coordinates.
2328
            ValueError: Raised if config["dataframe"]["time_stamp_alias"] not found
2329
                in Dataframe.
2330

2331
        Returns:
2332
            xr.DataArray: The computed normalization histogram (in TimeStamp units
2333
            per bin).
2334
        """
2335

2336
        if self._binned is None:
1✔
2337
            raise ValueError("Need to bin data first!")
1✔
2338
        if axis not in self._binned.coords:
1✔
2339
            raise ValueError(f"Axis '{axis}' not found in binned data!")
1✔
2340

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

2375
        return self._normalization_histogram
1✔
2376

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

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

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

2436
        input_types = map(type, [axes, bins, ranges])
1✔
2437
        allowed_types = [list, tuple]
1✔
2438

2439
        df = self._dataframe
1✔
2440

2441
        if not set(input_types).issubset(allowed_types):
1✔
2442
            raise TypeError(
×
2443
                "Inputs of axes, bins, ranges need to be list or tuple!",
2444
            )
2445

2446
        # Read out the values for the specified groups
2447
        group_dict_dd = {}
1✔
2448
        dfpart = df.get_partition(dfpid)
1✔
2449
        cols = dfpart.columns
1✔
2450
        for ax in axes:
1✔
2451
            group_dict_dd[ax] = dfpart.values[:, cols.get_loc(ax)]
1✔
2452
        group_dict = ddf.compute(group_dict_dd)[0]
1✔
2453

2454
        # Plot multiple histograms in a grid
2455
        grid_histogram(
1✔
2456
            group_dict,
2457
            ncol=ncol,
2458
            rvs=axes,
2459
            rvbins=bins,
2460
            rvranges=ranges,
2461
            backend=backend,
2462
            legend=legend,
2463
            histkwds=histkwds,
2464
            legkwds=legkwds,
2465
            **kwds,
2466
        )
2467

2468
    def save(
1✔
2469
        self,
2470
        faddr: str,
2471
        **kwds,
2472
    ):
2473
        """Saves the binned data to the provided path and filename.
2474

2475
        Args:
2476
            faddr (str): Path and name of the file to write. Its extension determines
2477
                the file type to write. Valid file types are:
2478

2479
                - "*.tiff", "*.tif": Saves a TIFF stack.
2480
                - "*.h5", "*.hdf5": Saves an HDF5 file.
2481
                - "*.nxs", "*.nexus": Saves a NeXus file.
2482

2483
            **kwds: Keyword argumens, which are passed to the writer functions:
2484
                For TIFF writing:
2485

2486
                - **alias_dict**: Dictionary of dimension aliases to use.
2487

2488
                For HDF5 writing:
2489

2490
                - **mode**: hdf5 read/write mode. Defaults to "w".
2491

2492
                For NeXus:
2493

2494
                - **reader**: Name of the nexustools reader to use.
2495
                  Defaults to config["nexus"]["reader"]
2496
                - **definiton**: NeXus application definition to use for saving.
2497
                  Must be supported by the used ``reader``. Defaults to
2498
                  config["nexus"]["definition"]
2499
                - **input_files**: A list of input files to pass to the reader.
2500
                  Defaults to config["nexus"]["input_files"]
2501
                - **eln_data**: An electronic-lab-notebook file in '.yaml' format
2502
                  to add to the list of files to pass to the reader.
2503
        """
2504
        if self._binned is None:
1✔
2505
            raise NameError("Need to bin data first!")
1✔
2506

2507
        if self._normalized is not None:
1✔
2508
            data = self._normalized
×
2509
        else:
2510
            data = self._binned
1✔
2511

2512
        extension = pathlib.Path(faddr).suffix
1✔
2513

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

2542
            if isinstance(input_files, str):
1✔
2543
                input_files = [input_files]
1✔
2544

2545
            if "eln_data" in kwds:
1✔
2546
                input_files.append(kwds.pop("eln_data"))
×
2547

2548
            to_nexus(
1✔
2549
                data=data,
2550
                faddr=faddr,
2551
                reader=reader,
2552
                definition=definition,
2553
                input_files=input_files,
2554
                **kwds,
2555
            )
2556

2557
        else:
2558
            raise NotImplementedError(
1✔
2559
                f"Unrecognized file format: {extension}.",
2560
            )
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

© 2025 Coveralls, Inc