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