chx-workflows/test_sparsify.py at main · NSLS2/chx-workflows · GitHub

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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
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
import filecmp
import numpy as np
import os
import pytest
import tiled

from dictdiffer import diff
from chx_compress.io.multifile.multifile import multifile_reader
from pandas import Timestamp
from pathlib import Path
from sparsify import get_metadata, sparsify, get_run_dask, get_run_sandbox_dask, get_tiled_client_sandbox_dask
from masks import MaskClient
from tiled.queries import Key
from prefect.testing.utilities import prefect_test_harness
from dotenv import load_dotenv

DATA_DIRECTORY = Path("/nsls2/data/chx/legacy/Compressed_Data")

mask_client = MaskClient(get_tiled_client_sandbox_dask())

run1 = get_run_dask("d85d157f-57d9-4649-9b65-0d3b9f754e01")
run2 = get_run_dask("e909f4a2-12e3-4521-a7a6-be2b728a826b")
run3 = get_run_dask("b79184e1-d053-42e4-b1eb-f8ab0a146220")

test_runs = (
    "14ed6885-2b6a-4645-85a2-a09413f618c2", # Just fixed the rootmap.
    "961b07ca-2133-46f9-8337-57053baa011b", # Very dense. Working.
    "eaf4d7df-2585-460a-8d9c-5a53613782e7", # Need the correct masks.
    "e909f4a2-12e3-4521-a7a6-be2b728a826b", # Working.
    "1bdfe080-250e-4874-b4cc-0b5acbdc99a3",
    "2e902219-5910-434f-9c89-75d7ea46db44",
    "4b31b18c-6207-491e-bca7-162431281a89",
    "5f1dc5a6-8d9f-479a-b2b5-5e2fabec594e",
    "6dff7417-6d41-4e71-ad3e-6982e838578f",
    "7d95f576-dfec-4bc7-905d-6b801420ca48",
    "81bcc380-d5cb-49f6-821d-4ca5ba0407e2",
    "8a8b73b8-5b8e-4dcd-81cd-c211520dccfb",
    "8af27b39-dd4c-4630-bdff-21af77b22227",
    "8ffddd5b-31f7-433f-babb-a2c995fd2587",
    "f9d3b730-1915-455e-985c-bb169fb74876",
)


def read_frame(multifile, image_index):
    image_array = np.zeros((multifile.header_info['ncols'],
                            multifile.header_info['nrows']))
    np.put(image_array, *multifile[image_index])
    return image_array.astype('uint16')


@pytest.fixture(autouse=True, scope="session")
def prefect_test_fixture():
    with prefect_test_harness():
        yield


@pytest.mark.parametrize("run_uid", test_runs)
def test_get_metadata(run_uid):
    """
    Check that the metadata from get_metadata matches the original
    metadata.
    """

    run = get_run_dask(run_uid)
    metadata_new = get_metadata(run)

    original_data = multifile_reader(
        f"{DATA_DIRECTORY}/uid_{run_uid}.cmp"
    ).header_info

    # TODO: Update chx_patches in tiled_site_config to get this metadata.
    exceptions = {
        "bytes",
        "byte_count",
        "rows_end",
        "nrows",
        "rows_begin",
        "cols_begin",
        "ncols",
        "cols_end",
    }

    assert metadata_original.keys() - metadata_new.keys() <= exceptions

    for key in metadata_original.keys() - exceptions:
        assert metadata_new[key] == metadata_original[key]


@pytest.mark.parametrize("run_uid", test_runs)
def test_sparsify(run_uid):
    """
    Make sure that the processed data from sparisfy
    matches the original proccesed data.
    """
    original_data = multifile_reader(
        f"{DATA_DIRECTORY}/uid_{run_uid}.cmp"
    )

    processed_uid = sparsify.fn(run_uid)
    new_data = get_run_sandbox_dask(processed_uid)

    for frame_number in range(new_data.shape[1]):
        assert np.array_equal(new_data[0][frame_number].todense(),
                              read_frame(original_data, frame_number))


@pytest.mark.parametrize("run_uid", test_runs)
def test_multifile(run_uid):
    """
    Make sure that the new multifile matches the original multifile.
    """
    original_file = f"{DATA_DIRECTORY}/uid_{run_uid}.cmp"
    processed_uid = sparsify.fn(run_uid)
    new_file = f"/tmp/{processed_uid}.cmp"
    processed_data = get_run_sandbox_dask(processed_uid)
    processed_data.export(new_file, format='application/x-eiger-multifile')
    assert filecmp.cmp(original_file, new_file)