pyrox.dev
1{
2 lib,
3 autograd,
4 buildPythonPackage,
5 fetchFromGitHub,
6 cvxopt,
7 cython,
8 jax,
9 jaxlib,
10 matplotlib,
11 numpy,
12 pymanopt,
13 pytestCheckHook,
14 pytest-cov-stub,
15 pythonOlder,
16 scikit-learn,
17 scipy,
18 setuptools,
19 tensorflow,
20 torch,
21}:
22
23buildPythonPackage rec {
24 pname = "pot";
25 version = "0.9.5";
26 pyproject = true;
27
28 disabled = pythonOlder "3.6";
29
30 src = fetchFromGitHub {
31 owner = "PythonOT";
32 repo = "POT";
33 tag = version;
34 hash = "sha256-sEK3uhZtjVJGEN1Gs8N0AMtiEOo9Kpn/zOSWUfGc/qE=";
35 };
36
37 build-system = [
38 setuptools
39 cython
40 numpy
41 ];
42
43 dependencies = [
44 numpy
45 scipy
46 ];
47
48 optional-dependencies = {
49 backend-numpy = [ ];
50 backend-jax = [
51 jax
52 jaxlib
53 ];
54 backend-cupy = [ ];
55 backend-tf = [ tensorflow ];
56 backend-torch = [ torch ];
57 cvxopt = [ cvxopt ];
58 dr = [
59 scikit-learn
60 pymanopt
61 autograd
62 ];
63 gnn = [
64 torch
65 # torch-geometric
66 ];
67 plot = [ matplotlib ];
68 all =
69 with optional-dependencies;
70 (
71 backend-numpy
72 ++ backend-jax
73 ++ backend-cupy
74 ++ backend-tf
75 ++ backend-torch
76 ++ optional-dependencies.cvxopt
77 ++ dr
78 ++ gnn
79 ++ plot
80 );
81 };
82
83 nativeCheckInputs = [
84 pytestCheckHook
85 pytest-cov-stub
86 ];
87
88 postPatch = ''
89 substituteInPlace setup.cfg \
90 --replace " --durations=20" "" \
91 --replace " --junit-xml=junit-results.xml" ""
92
93 # we don't need setup.py to find the macos sdk for us
94 sed -i '/sdk_path/d' setup.py
95 '';
96
97 # need to run the tests with the built package next to the test directory
98 preCheck = ''
99 pushd build/lib.*
100 ln -s -t . "$OLDPWD/test"
101 '';
102
103 postCheck = ''
104 popd
105 '';
106
107 disabledTests = [
108 # GPU tests are always skipped because of sandboxing
109 "warnings"
110 # Fixture is not available
111 "test_conditional_gradient"
112 "test_convert_between_backends"
113 "test_emd_backends"
114 "test_emd_emd2_types_devices"
115 "test_emd1d_type_devices"
116 "test_emd2_backends"
117 "test_factored_ot_backends"
118 "test_free_support_barycenter_backends"
119 "test_func_backends"
120 "test_generalized_conditional_gradient"
121 "test_line_search_armijo"
122 "test_loss_dual"
123 "test_max_sliced_backend"
124 "test_plan_dual"
125 "test_random_backends"
126 "test_sliced_backend"
127 "test_to_numpy"
128 "test_wasserstein_1d_type_devices"
129 "test_wasserstein"
130 "test_weak_ot_bakends"
131 # TypeError: Only integers, slices...
132 "test_emd1d_device_tf"
133 ];
134
135 pythonImportsCheck = [
136 "ot"
137 "ot.lp"
138 ];
139
140 meta = with lib; {
141 description = "Python Optimal Transport Library";
142 homepage = "https://pythonot.github.io/";
143 license = licenses.mit;
144 maintainers = with maintainers; [ yl3dy ];
145 };
146}