tjh.dev
1{
2 config,
3 lib,
4 stdenv,
5 buildPythonPackage,
6 fetchFromGitLab,
7
8 cudaSupport ? config.cudaSupport,
9
10 # build-system
11 setuptools,
12
13 # dependencies
14 astropy,
15 ducc0,
16 h5py,
17 jax,
18 jaxlib,
19 matplotlib,
20 mpi,
21 mpi4py,
22 numpy,
23 scipy,
24
25 # test
26 pytestCheckHook,
27 mpiCheckPhaseHook,
28 openssh,
29}:
30
31buildPythonPackage rec {
32 pname = "nifty8";
33 version = "8.5.2";
34 pyproject = true;
35
36 src = fetchFromGitLab {
37 domain = "gitlab.mpcdf.mpg.de";
38 owner = "ift";
39 repo = "nifty";
40 rev = "v${version}";
41 hash = "sha256-EWsJX+iqKOhQXEWlQfYUiPYqyfOfrwLtbI+DVn7vCQI=";
42 };
43
44 build-system = [ setuptools ];
45
46 dependencies = [
47 astropy
48 ducc0
49 h5py
50 jax
51 (jaxlib.override { inherit cudaSupport; })
52 matplotlib
53 mpi4py
54 mpi
55 numpy
56 scipy
57 ];
58
59 nativeCheckInputs = [
60 pytestCheckHook
61 mpiCheckPhaseHook
62 openssh
63 ];
64
65 checkPhase = ''
66 runHook preCheck
67
68 python3 -m pytest test
69
70 if [ "${stdenv.buildPlatform.system}" != "aarch64-linux" ] && \
71 [ "${stdenv.buildPlatform.system}" != "x86_64-darwin" ]; then
72 ${mpi}/bin/mpiexec -n 2 --bind-to none python3 -m pytest test/test_mpi
73 fi
74
75 runHook postCheck
76 '';
77
78 pythonImportsCheck = [ "nifty8" ];
79
80 meta = {
81 homepage = "https://gitlab.mpcdf.mpg.de/ift/nifty";
82 description = "Bayesian Imaging library for high-dimensional posteriors";
83 license = lib.licenses.gpl3;
84 maintainers = with lib.maintainers; [ parras ];
85 };
86}