1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5
6 # build-system
7 hatchling,
8
9 # dependencies
10 chex,
11 equinox,
12 jax,
13 jaxtyping,
14 optax,
15 scikit-learn,
16 tqdm,
17
18 # tests
19 pytestCheckHook,
20}:
21
22buildPythonPackage rec {
23 pname = "flowmc";
24 version = "0.4.5";
25 pyproject = true;
26
27 src = fetchFromGitHub {
28 owner = "kazewong";
29 repo = "flowMC";
30 tag = "flowMC-${version}";
31 hash = "sha256-D3K9cvmUvwsVAjvXdSDgYlqrzTYXVlSVQbfx7TANz8A=";
32 };
33
34 build-system = [ hatchling ];
35
36 pythonRelaxDeps = [
37 "jax"
38 ];
39
40 pythonRemoveDeps = [
41 # Not actual runtime dependencies
42 "pre-commit"
43 "pyright"
44 "pytest"
45 "ruff"
46 ];
47
48 dependencies = [
49 chex
50 equinox
51 jax
52 jaxtyping
53 optax
54 scikit-learn
55 tqdm
56 ];
57
58 pythonImportsCheck = [ "flowMC" ];
59
60 nativeCheckInputs = [ pytestCheckHook ];
61
62 meta = {
63 description = "Normalizing-flow enhanced sampling package for probabilistic inference in Jax";
64 homepage = "https://github.com/kazewong/flowMC";
65 changelog = "https://github.com/kazewong/flowMC/releases/tag/flowMC-${version}";
66 license = lib.licenses.mit;
67 maintainers = with lib.maintainers; [ GaetanLepage ];
68 };
69}