def test_multi_masks(lt_ctx): data = _mk_random(size=(16, 16, 16, 16), dtype="<u2") mask0 = _mk_random(size=(16, 16)) mask1 = sp.csr_matrix(_mk_random(size=(16, 16))) mask2 = sparse.COO.from_numpy(_mk_random(size=(16, 16))) expected = _naive_mask_apply([mask0, mask1, mask2], data) dataset = MemoryDataSet(data=data, tileshape=(4 * 4, 4, 4), num_partitions=2) analysis = lt_ctx.create_mask_analysis( dataset=dataset, factories=[lambda: mask0, lambda: mask1, lambda: mask2]) results = lt_ctx.run(analysis) assert np.allclose( results.mask_0.raw_data, expected[0], ) assert np.allclose( results.mask_1.raw_data, expected[1], ) assert np.allclose( results.mask_2.raw_data, expected[2], )
def test_start_local_cupyonly(hdf5_ds_1): cudas = detect()['cudas'] # Make sure we have enough partitions hdf5_ds_1.set_num_cores(len(cudas)) mask = _mk_random(size=(16, 16)) with hdf5_ds_1.get_reader().get_h5ds() as h5ds: data = h5ds[:] expected = _naive_mask_apply([mask], data) spec = cluster_spec(cpus=(), cudas=cudas, has_cupy=True) with DaskJobExecutor.make_local(spec=spec) as executor: ctx = api.Context(executor=executor) # Uses ApplyMasksUDF, which supports CuPy analysis = ctx.create_mask_analysis( dataset=hdf5_ds_1, factories=[lambda: mask] ) results = ctx.run(analysis) udf_res = ctx.run_udf(udf=DebugDeviceUDF(), dataset=hdf5_ds_1) # No CPU compute resources with pytest.raises(RuntimeError): _ = ctx.run_udf(udf=DebugDeviceUDF(backends=('numpy',)), dataset=hdf5_ds_1) cuda_res = ctx.run_udf(udf=DebugDeviceUDF(backends=('cuda',)), dataset=hdf5_ds_1) assert np.allclose( results.mask_0.raw_data, expected ) found = {} for val in udf_res['device_id'].data[0].values(): print(val) # no CPU assert val["cpu"] is None # Register which GPUs got work found[val["cuda"]] = True for val in cuda_res['device_id'].data[0].values(): print(val) # no CPU assert val["cpu"] is None # Register which GPUs got work found[val["cuda"]] = True for val in udf_res['backend'].data[0].values(): # use CuPy print(val) assert 'cupy' in val for val in cuda_res['backend'].data[0].values(): # no CuPy, i.e. NumPy print(val) assert 'numpy' in val # Test if each GPU got work. We have to see if this # actually works always since this depends on the scheduler behavior assert set(found.keys()) == set(cudas) assert np.all(udf_res['device_class'].data == 'cuda') assert np.allclose(udf_res['on_device'].data, data.sum(axis=(0, 1)))
def test_multi_mask_stack_force_dense(lt_ctx, backend): if backend == 'cupy': d = detect() cudas = detect()['cudas'] if not d['cudas'] or not d['has_cupy']: pytest.skip("No CUDA device or no CuPy, skipping CuPy test") try: if backend == 'cupy': set_use_cuda(cudas[0]) data = _mk_random(size=(16, 16, 16, 16), dtype="<u2") masks = sparse.COO.from_numpy(_mk_random(size=(2, 16, 16))) expected = _naive_mask_apply(masks, data) dataset = MemoryDataSet(data=data, tileshape=(4 * 4, 4, 4), num_partitions=2) analysis = lt_ctx.create_mask_analysis(dataset=dataset, factories=lambda: masks, use_sparse=False, mask_count=2) results = lt_ctx.run(analysis) assert np.allclose( results.mask_0.raw_data, expected[0], ) assert np.allclose( results.mask_1.raw_data, expected[1], ) finally: set_use_cpu(0)
def test_numerics(lt_ctx): dtype = 'float32' # Highest expected detector resolution RESOLUTION = 4096 # Highest expected detector dynamic range RANGE = 1e6 # default value for all cells # The test fails for 1.1 using float32! VAL = 1.0 data = np.full((2, 2, RESOLUTION, RESOLUTION), VAL, dtype=dtype) data[0, 0, 0, 0] += VAL * RANGE dataset = MemoryDataSet( data=data, tileshape=(2, RESOLUTION, RESOLUTION), num_partitions=2, sig_dims=2, ) mask0 = np.ones((RESOLUTION, RESOLUTION), dtype=dtype) analysis = lt_ctx.create_mask_analysis(dataset=dataset, factories=[lambda: mask0]) results = lt_ctx.run(analysis) expected = np.array( [[[VAL * RESOLUTION**2 + VAL * RANGE, VAL * RESOLUTION**2], [VAL * RESOLUTION**2, VAL * RESOLUTION**2]]]) naive = _naive_mask_apply([mask0], data) # print(expected) # print(naive) # print(results.mask_0.raw_data) assert np.allclose(expected, naive) assert np.allclose(expected[0], results.mask_0.raw_data)
def test_override_mask_dtype(lt_ctx): mask_dtype = np.float32 data = _mk_random(size=(16, 16, 16, 16), dtype=mask_dtype) masks = _mk_random(size=(2, 16, 16), dtype=np.float64) expected = _naive_mask_apply(masks.astype(mask_dtype), data) dataset = MemoryDataSet(data=data, tileshape=(4 * 4, 4, 4), num_partitions=2) analysis = lt_ctx.create_mask_analysis( dataset=dataset, factories=lambda: masks, mask_dtype=mask_dtype, mask_count=len(masks), ) results = lt_ctx.run(analysis) assert results.mask_0.raw_data.dtype == mask_dtype assert np.allclose( results.mask_0.raw_data, expected[0], ) assert np.allclose( results.mask_1.raw_data, expected[1], )
def test_multi_mask_force_dtype(lt_ctx): force_dtype = np.dtype(np.int32) data = _mk_random(size=(16, 16, 16, 16), dtype="int16") masks = _mk_random(size=(2, 16, 16), dtype="bool") expected = _naive_mask_apply(masks.astype(force_dtype), data.astype(force_dtype)) dataset = MemoryDataSet(data=data, tileshape=(4 * 4, 4, 4), num_partitions=2) analysis = lt_ctx.create_mask_analysis(dataset=dataset, factories=lambda: masks, dtype=force_dtype) results = lt_ctx.run(analysis) assert results.mask_0.raw_data.dtype.kind == force_dtype.kind assert results.mask_0.raw_data.dtype == force_dtype assert np.allclose( results.mask_0.raw_data, expected[0], ) assert np.allclose( results.mask_1.raw_data, expected[1], )
def test_ring_3d_ds(lt_ctx): data = _mk_random(size=(16 * 16, 16, 16)) dataset = MemoryDataSet( data=data.astype("<u2"), tileshape=(1, 16, 16), num_partitions=2, sig_dims=2, ) analysis = lt_ctx.create_ring_analysis(dataset=dataset, cx=8, cy=8, ri=5, ro=8) results = lt_ctx.run(analysis) mask = analysis.get_mask_factories()[0]() expected = _naive_mask_apply([mask], dataset.data.reshape( (16, 16, 16, 16))) assert results.intensity.raw_data.shape == (16 * 16, ) assert np.allclose( results.intensity.raw_data.reshape((16, 16)), expected, ) assert np.allclose( results.intensity_log.raw_data.reshape((16, 16)), expected, )
def test_numerics_succeed(lt_ctx): dtype = 'float64' # Highest expected detector resolution RESOLUTION = 4096 # Highest expected detector dynamic range RANGE = 1e6 # default value for all cells VAL = 1.1 data = np.full((2, 1, RESOLUTION, RESOLUTION), VAL, dtype=np.float32) data[0, 0, 0, 0] += VAL * RANGE dataset = MemoryDataSet( data=data, tileshape=(2, RESOLUTION, RESOLUTION), num_partitions=1, sig_dims=2, ) mask0 = np.ones((RESOLUTION, RESOLUTION), dtype=np.float32) analysis = lt_ctx.create_mask_analysis( dataset=dataset, factories=[lambda: mask0], mask_count=1, mask_dtype=dtype, ) results = lt_ctx.run(analysis) expected = np.array([[[VAL * RESOLUTION**2 + VAL * RANGE], [VAL * RESOLUTION**2]]]) naive = _naive_mask_apply([mask0.astype(dtype)], data.astype(dtype)) assert np.allclose(expected, naive) assert np.allclose(expected[0], results.mask_0.raw_data)
def test_endian(lt_ctx, TYPE): data = np.random.choice(a=0xFFFF, size=(16, 16, 16, 16)).astype(">u2") mask = _mk_random(size=(16, 16)) expected = _naive_mask_apply([mask], data) dataset = MemoryDataSet(data=data, tileshape=(4 * 4, 4, 4), num_partitions=2) _run_mask_test_program(lt_ctx, dataset, mask, expected, TYPE)
def test_mask_uint(lt_ctx, TYPE): data = _mk_random(size=(16, 16, 16, 16), dtype="<u2") mask = _mk_random(size=(16, 16)).astype("uint16") expected = _naive_mask_apply([mask], data) dataset = MemoryDataSet(data=data, tileshape=(4 * 4, 4, 4), num_partitions=2) _run_mask_test_program(lt_ctx, dataset, mask, expected, TYPE)
def test_subframe_tiles_fast(lt_ctx): data = _mk_random(size=(16, 16, 16, 16), dtype="<u2") mask = _mk_random(size=(16, 16)) expected = _naive_mask_apply([mask], data) dataset = MemoryDataSet(data=data, tileshape=(1, 8, 4, 4), partition_shape=(16, 16, 16, 16)) _run_mask_test_program(lt_ctx, dataset, mask, expected)
def test_signed(lt_ctx): data = np.random.choice(a=0xFFFF, size=(16, 16, 16, 16)).astype("<i4") mask = _mk_random(size=(16, 16)) expected = _naive_mask_apply([mask], data) dataset = MemoryDataSet(data=data, tileshape=(4, 4, 4, 4), partition_shape=(16, 16, 16, 16)) _run_mask_test_program(lt_ctx, dataset, mask, expected)
def test_single_frame_tiles(lt_ctx): data = _mk_random(size=(16, 16, 16, 16), dtype="<u2") mask = _mk_random(size=(16, 16)) expected = _naive_mask_apply([mask], data) dataset = MemoryDataSet(data=data, tileshape=(1, 16, 16), num_partitions=2) _run_mask_test_program(lt_ctx, dataset, mask, expected)
def test_ring_defaults(lt_ctx, ds_random): analysis = lt_ctx.create_ring_analysis(dataset=ds_random) results = lt_ctx.run(analysis) mask = analysis.get_mask_factories()[0]() expected = _naive_mask_apply([mask], ds_random.data) assert np.allclose( results.intensity.raw_data, expected, )
def test_point_complex(lt_ctx, ds_complex): analysis = lt_ctx.create_point_analysis(dataset=ds_complex) results = lt_ctx.run(analysis) mask = analysis.get_mask_factories()[0]() expected = _naive_mask_apply([mask], ds_complex.data) assert np.allclose( results.intensity_complex.raw_data, expected, )
def test_point_1(lt_ctx, ds_random): analysis = lt_ctx.create_point_analysis(dataset=ds_random, x=8, y=8) results = lt_ctx.run(analysis) mask = np.zeros((16, 16)) mask[8, 8] = 1 expected = _naive_mask_apply([mask], ds_random.data) assert np.allclose( results.intensity.raw_data, expected, )
def test_hdf5_3d_apply_masks(lt_ctx, hdf5_ds_3d): mask = _mk_random(size=(16, 16)) with hdf5_ds_3d.get_reader().get_h5ds() as h5ds: data = h5ds[:] expected = _naive_mask_apply([mask], data.reshape((1, 17, 16, 16))) analysis = lt_ctx.create_mask_analysis(dataset=hdf5_ds_3d, factories=[lambda: mask]) results = lt_ctx.run(analysis) assert np.allclose(results.mask_0.raw_data, expected)
def test_start_local_default(hdf5_ds_1, local_cluster_ctx): mask = _mk_random(size=(16, 16)) d = detect() cudas = d['cudas'] with hdf5_ds_1.get_reader().get_h5ds() as h5ds: data = h5ds[:] expected = _naive_mask_apply([mask], data) ctx = local_cluster_ctx analysis = ctx.create_mask_analysis(dataset=hdf5_ds_1, factories=[lambda: mask]) num_cores_ds = ctx.load('memory', data=np.zeros((2, 3, 4, 5))) workers = ctx.executor.get_available_workers() cpu_count = len(workers.has_cpu()) gpu_count = len(workers.has_cuda()) assert num_cores_ds._cores == max(cpu_count, gpu_count) # Based on ApplyMasksUDF, which is CuPy-enabled hybrid = ctx.run(analysis) _ = ctx.run_udf(udf=DebugDeviceUDF(), dataset=hdf5_ds_1) _ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy', 'numpy')), dataset=hdf5_ds_1) _ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cuda', 'numpy')), dataset=hdf5_ds_1) _ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy', 'cuda', 'numpy')), dataset=hdf5_ds_1) if cudas: cuda_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('cuda', 'numpy')), dataset=hdf5_ds_1, backends=('cuda', )) if d['has_cupy']: cupy_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy', 'numpy')), dataset=hdf5_ds_1, backends=('cupy', )) else: with pytest.raises(RuntimeError): cupy_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy', 'numpy')), dataset=hdf5_ds_1, backends=('cupy', )) cupy_only = None numpy_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('numpy', )), dataset=hdf5_ds_1) assert np.allclose(hybrid.mask_0.raw_data, expected) if cudas: assert np.all(cuda_only['device_class'].data == 'cuda') if cupy_only is not None: assert np.all(cupy_only['device_class'].data == 'cuda') assert np.all(numpy_only['device_class'].data == 'cpu')
def test_hdf5_apply_masks_1(lt_ctx, hdf5_ds_1, TYPE): mask = _mk_random(size=(16, 16)) with hdf5_ds_1.get_reader().get_h5ds() as h5ds: data = h5ds[:] expected = _naive_mask_apply([mask], data) analysis = lt_ctx.create_mask_analysis(dataset=hdf5_ds_1, factories=[lambda: mask]) analysis.TYPE = TYPE results = lt_ctx.run(analysis) assert np.allclose(results.mask_0.raw_data, expected)
def test_disk_defaults(lt_ctx, ds_random): analysis = lt_ctx.create_disk_analysis(dataset=ds_random) results = lt_ctx.run(analysis) mask = analysis.get_mask_factories()[0]() expected = _naive_mask_apply([mask], ds_random.data) # TODO: test the actual mask shape, too # at least mask application should match: assert np.allclose( results.intensity.raw_data, expected, )
def test_masks_complex_mask(lt_ctx, ds_complex): mask0 = _mk_random(size=(16, 16), dtype="complex64") analysis = lt_ctx.create_mask_analysis(dataset=ds_complex, factories=[lambda: mask0]) expected = _naive_mask_apply([mask0], ds_complex.data) results = lt_ctx.run(analysis) assert results.mask_0_complex.raw_data.shape == (16, 16) assert np.allclose(results.mask_0_complex.raw_data, expected) # also execute _run_mask_test_program to check sparse implementation. # _run_mask_test_program checks mask_0 result, which is np.abs(mask_0_complex) _run_mask_test_program(lt_ctx, ds_complex, mask0, np.abs(expected))
def test_weird_partition_shapes_1_slow(lt_ctx): data = _mk_random(size=(16, 16, 16, 16), dtype="<u2") mask = _mk_random(size=(16, 16)) expected = _naive_mask_apply([mask], data) dataset = MemoryDataSet(data=data, tileshape=(1, 16, 16), partition_shape=(16, 16, 2, 2)) _run_mask_test_program(lt_ctx, dataset, mask, expected) p = next(dataset.get_partitions()) t = next(p.get_tiles()) assert tuple(t.tile_slice.shape) == (1, 1, 2, 2)
def test_start_local(hdf5_ds_1): mask = _mk_random(size=(16, 16)) with hdf5_ds_1.get_reader().get_h5ds() as h5ds: data = h5ds[:] expected = _naive_mask_apply([mask], data) with api.Context() as ctx: analysis = ctx.create_mask_analysis(dataset=hdf5_ds_1, factories=[lambda: mask]) results = ctx.run(analysis) assert np.allclose(results.mask_0.raw_data, expected)
def test_numerics_fail(lt_ctx, TYPE): dtype = 'float32' # Highest expected detector resolution RESOLUTION = 4096 # Highest expected detector dynamic range RANGE = 1e6 # default value for all cells VAL = 1.1 data = np.full((2, 1, RESOLUTION, RESOLUTION), VAL, dtype=np.float32) data[0, 0, 0, 0] += VAL * RANGE dataset = MemoryDataSet( data=data, tileshape=(2, RESOLUTION, RESOLUTION), num_partitions=1, sig_dims=2, ) mask0 = np.ones((RESOLUTION, RESOLUTION), dtype=np.float64) analysis = lt_ctx.create_mask_analysis( dataset=dataset, factories=[lambda: mask0], mask_count=1, mask_dtype=dtype, ) analysis.TYPE = TYPE results = lt_ctx.run(analysis) expected = np.array([[[VAL * RESOLUTION**2 + VAL * RANGE], [VAL * RESOLUTION**2]]]) naive = _naive_mask_apply([mask0], data) naive_32 = _naive_mask_apply([mask0.astype(dtype)], data) # The masks are float64, that means the calculation is performed with high resolution # and the naive result should be correct assert np.allclose(expected, naive) # We make sure LiberTEM calculated this with the lower-precision dtype we set assert np.allclose(results.mask_0.raw_data, expected[0]) == np.allclose(naive_32, expected) # Confirm that the numerical precision is actually insufficient. # If this succeeds, we have to rethink the premise of this test. assert not np.allclose(results.mask_0.raw_data, expected[0])
def test_signed(lt_ctx, TYPE): data = np.random.choice(a=0xFFFF, size=(16, 16, 16, 16)).astype("<i4") mask = _mk_random(size=(16, 16)) expected = _naive_mask_apply([mask], data) # NOTE: we allow casting from int32 to float32 here, and may lose some # precision in case of data with large dynamic range dataset = MemoryDataSet( data=data, tileshape=(4 * 4, 4, 4), num_partitions=2, check_cast=False, ) _run_mask_test_program(lt_ctx, dataset, mask, expected, TYPE)
def test_ring_1(lt_ctx, ds_random): analysis = lt_ctx.create_ring_analysis(dataset=ds_random, cx=8, cy=8, ri=5, ro=8) results = lt_ctx.run(analysis) mask = analysis.get_mask_factories()[0]() expected = _naive_mask_apply([mask], ds_random.data) assert results.intensity.raw_data.shape == (16, 16) assert np.allclose( results.intensity.raw_data, expected, )
def test_weird_partition_shapes_1_fast(lt_ctx): # XXX MemoryDataSet is now using Partition3D and so on, so we can't create # partitions with weird shapes so easily anymore (in this case, partitioned in # the signal dimensions). maybe fix this with a custom DataSet impl that simulates this? data = _mk_random(size=(16, 16, 16, 16), dtype="<u2") mask = _mk_random(size=(16, 16)) expected = _naive_mask_apply([mask], data) dataset = MemoryDataSet(data=data, tileshape=(8, 16, 16), partition_shape=(16, 16, 8, 8)) _run_mask_test_program(lt_ctx, dataset, mask, expected) p = next(dataset.get_partitions()) t = next(p.get_tiles()) assert tuple(t.tile_slice.shape) == (1, 8, 8, 8)
def test_mask_job(lt_ctx): data = _mk_random(size=(16, 16, 16, 16), dtype="<u2") mask0 = _mk_random(size=(16, 16)) mask1 = sp.csr_matrix(_mk_random(size=(16, 16))) expected = _naive_mask_apply([mask0, mask1], data) dataset = MemoryDataSet(data=data, tileshape=(4, 4, 4, 4), partition_shape=(16, 16, 16, 16)) job = lt_ctx.create_mask_job( dataset=dataset, factories=[lambda: mask0, lambda: mask1] ) results = lt_ctx.run(job) assert np.allclose( results, expected, )
def test_mask_udf(lt_ctx): data = _mk_random(size=(16, 16, 16, 16), dtype="<u2") mask0 = _mk_random(size=(16, 16)) mask1 = sp.csr_matrix(_mk_random(size=(16, 16))) mask2 = sparse.COO.from_numpy(_mk_random(size=(16, 16))) # The ApplyMasksUDF returns data with shape ds.shape.nav + (mask_count, ), # different from ApplyMasksJob expected = np.moveaxis(_naive_mask_apply([mask0, mask1, mask2], data), (0, 1), (2, 0)) dataset = MemoryDataSet(data=data, tileshape=(4 * 4, 4, 4), num_partitions=2) udf = ApplyMasksUDF( mask_factories=[lambda: mask0, lambda: mask1, lambda: mask2] ) results = lt_ctx.run_udf(udf=udf, dataset=dataset) assert np.allclose(results['intensity'].data, expected)
def test_start_local_default(hdf5_ds_1): mask = _mk_random(size=(16, 16)) d = detect() cudas = d['cudas'] with hdf5_ds_1.get_reader().get_h5ds() as h5ds: data = h5ds[:] expected = _naive_mask_apply([mask], data) with api.Context() as ctx: analysis = ctx.create_mask_analysis(dataset=hdf5_ds_1, factories=[lambda: mask]) # Based on ApplyMasksUDF, which is CuPy-enabled hybrid = ctx.run(analysis) _ = ctx.run_udf(udf=DebugDeviceUDF(), dataset=hdf5_ds_1) _ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy', 'numpy')), dataset=hdf5_ds_1) _ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cuda', 'numpy')), dataset=hdf5_ds_1) _ = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy', 'cuda', 'numpy')), dataset=hdf5_ds_1) if cudas: cuda_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('cuda', 'numpy')), dataset=hdf5_ds_1, backends=('cuda', )) if d['has_cupy']: cupy_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('cupy', 'numpy')), dataset=hdf5_ds_1, backends=('cupy', )) else: with pytest.raises(RuntimeError): cupy_only = ctx.run_udf( udf=DebugDeviceUDF(backends=('cupy', 'numpy')), dataset=hdf5_ds_1, backends=('cupy', )) cupy_only = None numpy_only = ctx.run_udf(udf=DebugDeviceUDF(backends=('numpy', )), dataset=hdf5_ds_1) assert np.allclose(hybrid.mask_0.raw_data, expected) if cudas: assert np.all(cuda_only['device_class'].data == 'cuda') if cupy_only is not None: assert np.all(cupy_only['device_class'].data == 'cuda') assert np.all(numpy_only['device_class'].data == 'cpu')