utils.py 31.6 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
# coding: utf-8
# /*##########################################################################
#
# Copyright (c) 2019 European Synchrotron Radiation Facility
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# ###########################################################################*/
"""An :class:`.Enum` class with additional features."""

from __future__ import absolute_import

payno's avatar
payno committed
29
__authors__ = ["T. Vincent", "H. Payno"]
30
31
32
33
34
__license__ = "MIT"
__date__ = "29/04/2019"


import enum
35
36
import sys
import typing
37
38
import numpy
import os
39
import contextlib
40
from tomoscan.io import HDF5File
payno's avatar
payno committed
41
from tomoscan.esrf.hdf5scan import HDF5TomoScan
42
from silx.io.url import DataUrl
43
from silx.io.utils import get_data
44
45
import logging
import h5py
46
from silx.utils.deprecation import deprecated
47
48
49
from nxtomomill.converter.hdf5.acquisition.baseacquisition import (
    DatasetReader,
)
Tomas Farago's avatar
Tomas Farago committed
50
51
52
53
54

try:
    import hdf5plugin
except ImportError:
    pass
55
from collections.abc import Iterable
56
from silx.utils.enum import Enum as _Enum
57
import uuid
58
from silx.io.utils import h5py_read_dataset
59
60
import h5py._hl.selections as selection
from h5py import h5s as h5py_h5s
61
62
63
64
65
66
67
68
69
70
71
72


class ImageKey(_Enum):
    """
    Possible values of image_key_control
    """

    ALIGNMENT = -1
    PROJECTION = 0
    FLAT_FIELD = 1
    DARK_FIELD = 2
    INVALID = 3
73
74


75
76
77
78
79
class FieldOfView(_Enum):
    FULL = "Full"
    HALF = "Half"


80
81
@contextlib.contextmanager
def cwd_context():
payno's avatar
payno committed
82
83
84
85
86
    curdir = os.getcwd()
    try:
        yield
    finally:
        os.chdir(curdir)
87
88


89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
class Enum(enum.Enum):
    """Enum with additional class methods."""

    @classmethod
    def from_value(cls, value):
        """Convert a value to corresponding Enum member

        :param value: The value to compare to Enum members
           If it is already a member of Enum, it is returned directly.
        :return: The corresponding enum member
        :rtype: Enum
        :raise ValueError: In case the conversion is not possible
        """
        if isinstance(value, cls):
            return value
        for member in cls:
            if value == member.value:
                return member
        raise ValueError("Cannot convert: %s" % value)

    @classmethod
    def members(cls):
        """Returns a tuple of all members.

        :rtype: Tuple[Enum]
        """
        return tuple(member for member in cls)

    @classmethod
    def names(cls):
        """Returns a tuple of all member names.

        :rtype: Tuple[str]
        """
        return tuple(member.name for member in cls)

    @classmethod
    def values(cls):
        """Returns a tuple of all member values.

        :rtype: Tuple
        """
        return tuple(member.value for member in cls)


134
def embed_url(url: DataUrl, output_file: str) -> DataUrl:
135
136
137
    """
    Create a dataset under duplicate_data and with a random name
    to store it
138
139
140
141
142

    :param DataUrl url: dataset to be copied
    :param str output_file: where to store the dataset
    :param expected_type: some metadata to put in copied dataset attributes
    :type expected_type: Union[None, AcquisitionStep]
143
144
    :param data: data loaded from url is already loaded
    :type data: None, numpy.array
145
146
147
148
149
150
151
152
153
154
    """
    if not isinstance(url, DataUrl):
        return url
    elif url.file_path() == output_file:
        return url
    else:
        embed_data_path = "/".join(("/duplicate_data", str(uuid.uuid1())))
        with cwd_context():
            os.chdir(os.path.dirname(os.path.abspath(output_file)))
            with HDF5File(output_file, "a") as h5s:
155
                h5s[embed_data_path] = get_data(url)
156
                h5s[embed_data_path].attrs["original_url"] = url.path()
157
158
159
160
161
            return DataUrl(
                file_path=output_file, data_path=embed_data_path, scheme="silx"
            )


162
class FileExtension(Enum):
payno's avatar
payno committed
163
164
165
    H5 = ".h5"
    HDF5 = ".hdf5"
    NX = ".nx"
166
167


payno's avatar
payno committed
168
class Format(Enum):
169
    STANDARD = "standard"
payno's avatar
payno committed
170
    XRD_CT = "xrd-ct"
171
    XRD_3D = "3d-xrd"
payno's avatar
payno committed
172
173


174
175
176
177
178
179
180
181
182
183
184
185
186
187
def get_file_name(file_name, extension, check=True):
    """
    set the given extension

    :param str file_name: name of the file
    :param str extension: extension to give
    :param bool check: if check, already check if the file as one of the
                       '_FileExtension'
    """
    extension = FileExtension.from_value(extension.lower())
    if check:
        for value in FileExtension.values():
            if file_name.lower().endswith(value):
                return file_name
188
189
190
191
192
    return file_name + extension.value()


class Progress:
    """Simple interface for defining advancement on a 100 percentage base"""
payno's avatar
payno committed
193

194
    def __init__(self, name: str):
195
196
197
198
        self._name = name
        self.set_name(name)

    def set_name(self, name):
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
        self._name = name
        self.reset()

    def reset(self, max_: typing.Union[None, int] = None) -> None:
        """
        reset the advancement to n and max advancement to max_
        :param int max_:
        """
        self._n_processed = 0
        self._max_processed = max_

    def start_process(self) -> None:
        self.set_advancement(0)

    def set_advancement(self, value: int) -> None:
        """

        :param int value: set advancement to value
        """
        length = 20  # modify this to change the length
payno's avatar
payno committed
219
        block = int(round(length * value / 100))
220
        blocks_str = "#" * block + "-" * (length - block)
payno's avatar
payno committed
221
        msg = "\r{0}: [{1}] {2}%".format(self._name, blocks_str, round(value, 2))
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
        if value >= 100:
            msg += " DONE\r\n"
        sys.stdout.write(msg)
        sys.stdout.flush()

    def end_process(self) -> None:
        """Set advancement to 100 %"""
        self.set_advancement(100)

    def set_max_advancement(self, n: int) -> None:
        """

        :param int n: number of steps contained by the advancement. When
        advancement reach this value, advancement will be 100 %
        """
        self._max_processed = n

    def increase_advancement(self, i: int = 1) -> None:
        """

        :param int i: increase the advancement of n step
        """
        self._n_processed += i
245
        advancement = int(float(self._n_processed / self._max_processed) * 100)
246
247
        self.set_advancement(advancement)

248
249
250

def get_tuple_of_keys_from_cmd(cmd_value: str) -> tuple:
    """Return a tuple"""
payno's avatar
payno committed
251
    return tuple(cmd_value.split(","))
252
253


payno's avatar
payno committed
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
def is_nx_tomo_entry(file_path, entry):
    """

    :param str file_path: hdf5 file path
    :param str entry: entry to check
    :return: True if the entry is an NXTomo entry
    """
    if not os.path.exists(file_path):
        return False
    else:
        with HDF5File(file_path, mode="r") as h5s:
            if entry not in h5s:
                return False
            node = h5s[entry]
            return HDF5TomoScan.node_is_nxtomo(node)


271
def add_dark_flat_nx_file(
272
    file_path: str,
273
    entry: str,
payno's avatar
payno committed
274
275
276
277
    darks_start: typing.Union[None, numpy.ndarray, DataUrl] = None,
    flats_start: typing.Union[None, numpy.ndarray, DataUrl] = None,
    darks_end: typing.Union[None, numpy.ndarray, DataUrl] = None,
    flats_end: typing.Union[None, numpy.ndarray, DataUrl] = None,
278
    extras: typing.Union[None, dict] = None,
279
    logger: typing.Union[None, logging.Logger] = None,
280
    embed_data: bool = False,
281
282
283
284
):
    """
    This will get all data from entry@input_file and patch them with provided
    dark and / or flat(s).
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
    We consider the sequence as: dark, start_flat, projections, end_flat.

    Behavior regarding data type and target dataset:

    * if dataset at `entry` already exists:
        * if dataset at `entry` is a 'standard' dataset:
            * data will be loaded if necessary and `enrty` will be updated
        * if dataset at `entry` is a virtual dataset:
            * if `data` is a numpy array then we raise an error: the data should
              already be saved somewhere and you should provide a DataUrl
            * if `data` is a DataUrl then the virtual dataset is updated and
              a virtual source pointing to the
              DataUrl.file_path()@DataUrl.data_path() is added to the layout
    * if a new dataset `entry` need to be added:
        * if `data` is a numpy array then we create a new 'standard' Dataset
        * if `data` is a DataUrl then a new virtual dataset will be created

302
303
304
    note: Datasets `image_key`, `image_key_control`, `rotation_angle` and
    `count_time` will be copied each time.

305
306
    :param file_path: NXTomo file containing data to be patched
    :type file_path: str
307
308
    :param entry: entry to be patched
    :type entry: str
payno's avatar
payno committed
309
310
311
312
313
314
315
316
317
    :param darks_start: (3D) numpy array containing the first dark serie if any
    :type darks_start: Union[None, numpy.ndarray, DataUrl]
    :param flats_start: (3D) numpy array containing the first flat if any
    :type flats_start: Union[None, numpy.ndarray, DataUrl]
    :param darks_end: (3D) numpy array containing dark the second dark serie if
                      any
    :type darks_end: Union[None, numpy.ndarray, DataUrl]
    :param flats_end: (3D) numpy array containing the second flat if any
    :type flats_end: Union[None, numpy.ndarray, DataUrl]
318
319
    :param extras: dictionary to specify some parameters for flats and dark
                   like rotation angle.
payno's avatar
payno committed
320
321
                   valid keys: 'start_dark', 'end_dark', 'start_flag',
                   'end_flag'.
322
323
324
                   Values should be a dictionary of 'NXTomo' keys with
                   values to be set instead of 'default values'.
                   Possible values are:
325
326
                   * `count_time`
                   * `rotation_angle`
327
    :type extras: Union[None, dict]
328
    :param Union[None, logging.Logger] logger: object for logs
329
330
    :param bool embed_data: if True then each external data will be copy
                            under a 'duplicate_data' folder
331
    """
332
333
    if extras is None:
        extras = {}
payno's avatar
payno committed
334
335
336
337
338
339
340
341
    else:
        for key in extras:
            valid_extra_keys = ("darks_start", "darks_end", "flats_start", "flats_end")
            if key not in valid_extra_keys:
                raise ValueError(
                    "{key} is not recognized. Valid values are "
                    "{keys}".format(key=key, keys=valid_extra_keys)
                )
342

343
    if embed_data is True:
344
345
346
347
        darks_start = embed_url(darks_start, output_file=file_path)
        darks_end = embed_url(darks_end, output_file=file_path)
        flats_start = embed_url(flats_start, output_file=file_path)
        flats_end = embed_url(flats_end, output_file=file_path)
348
349
350
351
352
353
354
355
356
357
358
    else:
        for url in (darks_start, darks_end, flats_start, flats_end):
            if url is not None and isinstance(url, DataUrl):
                if isinstance(url.data_slice(), slice):
                    if url.data_slice().step not in (None, 1):
                        raise ValueError(
                            "When data is not embed slice `step`"
                            "must be None or 1. Other values are"
                            "not handled. Failing url is {}"
                            "".format(url)
                        )
359

360
    # !!! warning: order of dark / flat treatments import
payno's avatar
payno committed
361
362
    data_names = "flats_start", "darks_end", "flats_end", "darks_start"
    datas = flats_start, darks_end, flats_end, darks_start
363
364
    keys_value = (
        ImageKey.FLAT_FIELD.value,
payno's avatar
payno committed
365
        ImageKey.DARK_FIELD.value,
366
367
368
        ImageKey.FLAT_FIELD.value,
        ImageKey.DARK_FIELD.value,
    )
payno's avatar
payno committed
369
    wheres = "start", "end", "end", "start"  # warning: order import
370

371
372
373
374
375
376
    for d_n, data, key, where in zip(data_names, datas, keys_value, wheres):
        if data is None:
            continue
        n_frames_to_insert = 1
        if isinstance(data, numpy.ndarray) and data.ndim == 3:
            n_frames_to_insert = data.shape[0]
377
378
379
380
381
382
383
384
        elif isinstance(data, DataUrl):
            with HDF5File(data.file_path(), mode="r") as h5s:
                if data.data_path() not in h5s:
                    raise KeyError(
                        "Path given ({}) is not in {}".format(
                            data.data_path(), data.file_path
                        )
                    )
payno's avatar
payno committed
385
                data_node = get_data(data)
payno's avatar
payno committed
386
                if data_node.ndim == 3:
387
                    n_frames_to_insert = data_node.shape[0]
388
        else:
389
            raise TypeError("{} as input is not managed".format(type(data)))
390
391
392
393
394
395
396
397

        if logger is not None:
            logger.info(
                "insert {data_type} frame of type {key} at the"
                "{where}".format(data_type=type(data), key=key, where=where)
            )
        # update 'data' dataset
        data_path = os.path.join(entry, "instrument", "detector", "data")
398
        _FrameAppender(
399
            data, file_path, data_path=data_path, where=where, logger=logger
400
        ).process()
401
402
403
404
405
        # update image-key and image_key_control (we are not managing the
        # 'alignment projection here so values are identical')
        ik_path = os.path.join(entry, "instrument", "detector", "image_key")
        ikc_path = os.path.join(entry, "instrument", "detector", "image_key_control")
        for path in (ik_path, ikc_path):
406
            _FrameAppender(
407
408
409
410
411
                [key] * n_frames_to_insert,
                file_path,
                data_path=path,
                where=where,
                logger=logger,
412
            ).process()
413
414
415
416
417
418
419
420

        # add 'other' necessaries key:
        count_time_path = os.path.join(entry, "instrument", "detector", "count_time")
        rotation_angle_path = os.path.join(entry, "sample", "rotation_angle")
        data_key_paths = (count_time_path, rotation_angle_path)
        data_keys = ("count_time", "rotation_angle")

        for data_key, data_key_path in zip(data_keys, data_key_paths):
421
422
            if d_n in extras and data_key in extras[d_n]:
                provided_value = extras[d_n][data_key]
423
424
425
426
427
428
429
430
431
432
433
434
435
436
                if isinstance(provided_value, Iterable):
                    if len(provided_value) != n_frames_to_insert:
                        raise ValueError(
                            "Given value to store from extras has"
                            " incoherent length({}) compare to "
                            "the number of frame to save ({})"
                            "".format(len(provided_value), n_frames_to_insert)
                        )
                    else:
                        data_to_insert = provided_value
                else:
                    try:
                        data_to_insert = [provided_value] * n_frames_to_insert
                    except Exception as e:
payno's avatar
payno committed
437
                        logger.error("Fail to create data to insert. Error is", e)
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
                        return
            else:
                # get default values
                def get_default_value(location, where_):
                    with HDF5File(file_path, mode="r") as h5s:
                        if location not in h5s:
                            return None
                        existing_data = h5s[location]
                        if where_ == "start":
                            return existing_data[0]
                        else:
                            return existing_data[-1]

                default_value = get_default_value(data_key_path, where)
                if default_value is None:
                    raise ValueError(
                        "Unable to define a default value for "
                        "{}. Location empty in {}."
                        "".format(data_key_path, file_path)
                    )
                elif logger:
                    logger.debug(
                        "No value(s) provided for {path}. "
                        "Extract some default value ({def_value})."
                        "".format(path=data_key_path, def_value=default_value)
                    )
                data_to_insert = [default_value] * n_frames_to_insert

466
            _FrameAppender(
467
                data_to_insert,
468
                file_path,
469
                data_path=data_key_path,
470
471
                where=where,
                logger=logger,
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
            ).process()


class _FrameAppender:
    """
    Class to insert 2D frame(s) to an existing dataset
    """

    def __init__(self, data, file_path, data_path, where, logger=None):
        if where not in ("start", "end"):
            raise ValueError("`where` should be `start` or `end`")

        if not isinstance(data, (DataUrl, numpy.ndarray, list, tuple)):
            raise TypeError(
                "data should be an instance of DataUrl or a numpy "
                "array not {}".format(type(data))
488
489
            )

490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
        self.data = data
        self.file_path = file_path
        self.data_path = data_path
        self.where = where
        self.logger = logger

    def process(self) -> None:
        """
        main function. Will start the insertion of frame(s)
        """
        with HDF5File(self.file_path, mode="a") as h5s:
            if self.data_path in h5s:
                self._add_to_existing_dataset(h5s)
            else:
                self._create_new_dataset(h5s)
            if self.logger:
                self.logger.info(
                    "data added to {entry}@{file_path}"
                    "".format(entry=self.data_path, file_path=self.file_path)
                )
510

511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
    def _add_to_existing_virtual_dataset(self, h5s):
        if (
            h5py.version.hdf5_version_tuple[0] <= 1
            and h5py.version.hdf5_version_tuple[1] < 12
        ):
            if self.logger:
                self.logger.warning(
                    "You are working on virtual dataset"
                    "with a hdf5 version < 12. Frame "
                    "you want to change might be "
                    "modified depending on the working "
                    "directory without notifying."
                    "See https://github.com/silx-kit/silx/issues/3277"
                )
        if isinstance(self.data, (numpy.ndarray, list, tuple)):
            raise TypeError(
                "Provided data is a numpy array when given"
                "dataset path is a virtual dataset. "
                "You must store the data somewhere else "
                "and provide a DataUrl"
            )
        else:
            if self.logger is not None:
                self.logger.debug(
                    "Update virtual dataset: "
                    "{entry}@{file_path}"
                    "".format(entry=self.data_path, file_path=self.file_path)
                )
            # store DataUrl in the current virtual dataset
            url = self.data
541

542
543
544
545
546
            data_need_reshape = False
            with DatasetReader(self.data) as data_frm_url:
                if not data_frm_url.ndim in (2, 3):
                    raise ValueError(
                        "{} should point to 2D or 3D dataset ".format(url.path())
547
                    )
548
549
550
551
552
553
554
555
                if data_frm_url.ndim == 2:
                    new_shape = 1, data_frm_url.shape[0], data_frm_url.shape[1]
                    if self.logger is not None:
                        self.logger.info(
                            "reshape provided data to 3D (from {} to {})"
                            "".format(data_frm_url.shape, new_shape)
                        )
                    data_need_reshape = True
556

557
558
559
            if data_need_reshape:
                data_frm_url = get_data(data_frm_url)
                self.__insert_in_vds(h5s, url, data_frm_url)
560
            else:
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
                with DatasetReader(self.data) as data_frm_url:
                    self.__insert_in_vds(h5s, url, data_frm_url)

    def __insert_in_vds(self, h5s, url, data_frm_url):
        n_frames, dim_2, dim_1 = data_frm_url.shape

        virtual_sources_len = []
        virtual_sources = []
        # we need to recreate the VirtualSource they are not
        # store or available from the API
        for vs_info in h5s[self.data_path].virtual_sources():
            length, vs = self._recreate_vs(vs_info=vs_info, vds_file=self.file_path)
            virtual_sources.append(vs)
            virtual_sources_len.append(length)

        vds_file_path = os.path.abspath(os.path.relpath(url.file_path(), os.getcwd()))
        vds_file_path = os.path.realpath(vds_file_path)
        vds_file_path = os.path.relpath(vds_file_path, os.path.dirname(self.file_path))
        if not vds_file_path.startswith("./"):
            vds_file_path = "./" + vds_file_path

        new_virtual_source = h5py.VirtualSource(
            path_or_dataset=vds_file_path,
            name=url.data_path(),
            shape=data_frm_url.shape,
        )

        if url.data_slice() is not None:
            # in the case we have to process to a FancySelection
            with HDF5File(os.path.abspath(url.file_path()), mode="r") as h5sd:
                dst = h5sd[url.data_path()]
                sel = selection.select(
                    h5sd[url.data_path()].shape, url.data_slice(), dst
                )
                new_virtual_source.sel = sel

        n_frames += h5s[self.data_path].shape[0]
        data_type = h5s[self.data_path].dtype

        if self.where == "start":
            virtual_sources.insert(0, new_virtual_source)
            virtual_sources_len.insert(0, data_frm_url.shape[0])
        else:
            virtual_sources.append(new_virtual_source)
            virtual_sources_len.append(data_frm_url.shape[0])

        # create the new virtual dataset
        layout = h5py.VirtualLayout(shape=(n_frames, dim_2, dim_1), dtype=data_type)
        last = 0
        for v_source, vs_len in zip(virtual_sources, virtual_sources_len):
            layout[last : vs_len + last] = v_source
            last += vs_len
        if self.data_path in h5s:
            del h5s[self.data_path]
        h5s.create_virtual_dataset(self.data_path, layout)
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703

    def _add_to_existing_none_virtual_dataset(self, h5s):
        """
        for now when we want to add data *to a none virtual dataset*
        we always duplicate data if provided from a DataUrl.
        We could create a virtual dataset as well but seems to complicated for
        a use case that we don't really have at the moment.

        :param h5s:
        """
        if self.logger is not None:
            self.logger.debug("Update dataset: {entry}@{file_path}")
        if isinstance(self.data, (numpy.ndarray, list, tuple)):
            new_data = self.data
        else:
            url = self.data
            new_data = get_data(url)

        if isinstance(new_data, numpy.ndarray):
            if not new_data.shape[1:] == h5s[self.data_path].shape[1:]:
                raise ValueError(
                    "Data shapes are incoherent: {} vs {}".format(
                        new_data.shape, h5s[self.data_path].shape
                    )
                )

            new_shape = (
                new_data.shape[0] + h5s[self.data_path].shape[0],
                new_data.shape[1],
                new_data.shape[2],
            )
            data_to_store = numpy.empty(new_shape)
            if self.where == "start":
                data_to_store[: new_data.shape[0]] = new_data
                data_to_store[new_data.shape[0] :] = h5py_read_dataset(
                    h5s[self.data_path]
                )
            else:
                data_to_store[: h5s[self.data_path].shape[0]] = h5py_read_dataset(
                    h5s[self.data_path]
                )
                data_to_store[h5s[self.data_path].shape[0] :] = new_data
        else:
            assert isinstance(
                self.data, (list, tuple)
            ), "Unmanaged data type {}".format(type(self.data))
            o_data = h5s[self.data_path]
            o_data = list(h5py_read_dataset(o_data))
            if self.where == "start":
                new_data.extend(o_data)
                data_to_store = numpy.asarray(new_data)
            else:
                o_data.extend(new_data)
                data_to_store = numpy.asarray(o_data)

        del h5s[self.data_path]
        h5s[self.data_path] = data_to_store

    def _add_to_existing_dataset(self, h5s):
        """Add the frame to an existing dataset"""
        if h5s[self.data_path].is_virtual:
            self._add_to_existing_virtual_dataset(h5s=h5s)
        else:
            self._add_to_existing_none_virtual_dataset(h5s=h5s)

    def _create_new_dataset(self, h5s):
        """
        needs to create a new dataset. In this case the policy is:
           - if a DataUrl is provided then we create a virtual dataset
           - if a numpy array is provided then we create a 'standard' dataset
        """

        if isinstance(self.data, DataUrl):
            url = self.data

            with HDF5File(url.file_path(), mode="r") as o_h5s:
                if not o_h5s[url.data_path()].ndim in (2, 3):
                    raise ValueError(
                        "{} should point to 2D or 3D dataset ".format(url.path())
                    )
                data_shape = o_h5s[url.data_path()].shape
                data_type = o_h5s[url.data_path()].dtype
                if len(data_shape) == 2:
                    data_shape = (1, data_shape[0], data_shape[1])
                layout = h5py.VirtualLayout(shape=data_shape, dtype=data_type)
                layout[:] = h5py.VirtualSource(
                    url.file_path(), url.data_path(), shape=data_shape
                )
704
                h5s.create_virtual_dataset(self.data_path, layout)
705
706
707
708
        else:
            h5s[self.data_path] = self.data

    @staticmethod
709
    def _recreate_vs(vs_info, vds_file):
710
        """Simple util to retrieve a h5py.VirtualSource from virtual source
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
        information.

        to understand clearly this function you might first have a look at
        the use case exposed in issue:
        https://gitlab.esrf.fr/tomotools/nxtomomill/-/issues/40
        """
        with cwd_context():
            if os.path.dirname(vds_file) not in ("", None):
                os.chdir(os.path.dirname(vds_file))
            with HDF5File(vs_info.file_name, mode="r") as vs_node:
                dataset = vs_node[vs_info.dset_name]
                select_bounds = vs_info.vspace.get_select_bounds()
                left_bound = select_bounds[0]
                right_bound = select_bounds[1]
                length = right_bound[0] - left_bound[0] + 1
                # warning: for now step is not managed with virtual
                # dataset

                virtual_source = h5py.VirtualSource(
                    vs_info.file_name,
                    vs_info.dset_name,
                    shape=dataset.shape,
                )
                # here we could provide dataset but we won't to
                # insure file path will be relative.
                type_code = vs_info.src_space.get_select_type()
                # check for unlimited selections in case where selection is regular
                # hyperslab, which is the only allowed case for h5s.UNLIMITED to be
                # in the selection
                if (
payno's avatar
payno committed
741
742
                    type_code == h5py_h5s.SEL_HYPERSLABS
                    and vs_info.src_space.is_regular_hyperslab()
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
                ):
                    (
                        source_start,
                        stride,
                        count,
                        block,
                    ) = vs_info.src_space.get_regular_hyperslab()
                    source_end = source_start[0] + length

                    sel = selection.select(
                        dataset.shape,
                        slice(source_start[0], source_end),
                        dataset=dataset,
                    )
                    virtual_source.sel = sel

759
760
761
762
                return (
                    length,
                    virtual_source,
                )
763
764
765


@deprecated(replacement="_FrameAppender", since_version="0.5.0")
766
def _insert_frame_data(data, file_path, data_path, where, logger=None):
767
768
769
770
771
772
773
774
775
776
777
778
779
    """
    This function is used to insert some frame(s) (numpy 2D or 3D to an
    existing dataset. Before the existing array or After.

    :param data:
    :param file_path:
    :param data_path: If the path point to a virtual dataset them this one
                      will be updated but data should be a DataUrl. Of the
                      same shape. Else we will update the data_path by
                      extending the dataset.
    :param where:
    :raises TypeError: In the case the data type and existing data_path are
                       incompatible.
780
    """
781
782
783
784
    fa = _FrameAppender(
        data=data, file_path=file_path, data_path=data_path, where=where, logger=logger
    )
    return fa.process()
785
786
787
788
789
790
791


def change_image_key_control(
    file_path: str,
    entry: str,
    frames_indexes: typing.Union[slice, Iterable],
    image_key_control_value: typing.Union[int, ImageKey],
792
    logger=None,
793
794
795
796
797
798
799
800
801
802
803
804
):
    """
    Will modify image_key and image_key_control values for the requested
    frames.

    :param str file_path: path the nexus file
    :param str entry: name of the entry to modify
    :param frames_indexes: index of the frame for which we want to modify
                           the image key
    :type frames_indexes: Union[slice, Iterable]
    :param image_key_control_value:
    :type image_key_control_value: Union[int, ImageKey]
805
    :param logging.Logger logger: logger
806
807
    """
    if not isinstance(frames_indexes, (Iterable, slice)):
payno's avatar
payno committed
808
        raise TypeError("`frame_indexes` should be an instance of Iterable slice")
809
810
811
812
813
814
815
816
817
818
819
820
    if logger:
        logger.info(
            "Update frames {frames_indexes} to"
            "{image_key_control_value} of {entry}@{file_path}"
            "".format(
                frames_indexes=frames_indexes,
                image_key_control_value=image_key_control_value,
                entry=entry,
                file_path=file_path,
            )
        )

821
822
823
824
    image_key_control_value = ImageKey.from_value(image_key_control_value)
    with HDF5File(file_path, mode="a") as h5s:
        node = h5s[entry]
        image_keys_path = "/".join(("instrument", "detector", "image_key"))
825
        image_keys = h5py_read_dataset(node[image_keys_path])
826
827
828
        image_keys_control_path = "/".join(
            ("instrument", "detector", "image_key_control")
        )
829
        image_keys_control = h5py_read_dataset(node[image_keys_control_path])
830
        # filter frame indexes
831
        if isinstance(frames_indexes, slice):
832
833
834
835
            step = frames_indexes.step
            if step is None:
                step = 1
            stop = frames_indexes.stop
836
            if stop in (None, -1):
837
838
                stop = len(image_keys)
            frames_indexes = list(range(frames_indexes.start, stop, step))
839
        frames_indexes = list(
840
            filter(lambda x: 0 <= x <= len(image_keys_control), frames_indexes)
841
        )
842
843
844
845
846
847
848
849
850
851
        # manage image_key_control
        image_keys_control[frames_indexes] = image_key_control_value.value
        node[image_keys_control_path][:] = image_keys_control
        # manage image_key. In this case we should get rid of Alignment values
        # and replace it by Projection values
        image_key_value = image_key_control_value
        if image_key_value is ImageKey.ALIGNMENT:
            image_key_value = ImageKey.PROJECTION
        image_keys[frames_indexes] = image_key_value.value
        node[image_keys_path][:] = image_keys