Sirius python module (#338395)

pkgs/by-name/si/sirius/package.nix

@@ -21,6 +21,8 @@
 , boost
 , eigen
 , libvdwxc
+, enablePython ? false
+, pythonPackages ? null
 , llvmPackages
 , cudaPackages
 , rocmPackages
@@ -35,6 +37,7 @@
 }:
 
 assert builtins.elem gpuBackend [ "none" "cuda" "rocm" ];
+assert enablePython -> pythonPackages != null;
 
 stdenv.mkDerivation rec {
   pname = "SIRIUS";
@@ -81,10 +84,23 @@ stdenv.mkDerivation rec {
   ] ++ lib.optionals (gpuBackend == "rocm") [
     rocmPackages.clr
     rocmPackages.rocblas
-  ] ++ lib.optional stdenv.isDarwin llvmPackages.openmp
-  ;
+  ] ++ lib.optionals stdenv.isDarwin [
+    llvmPackages.openmp
+  ] ++ lib.optionals enablePython (with pythonPackages; [
+    python
+    pybind11
+  ]);
 
-  propagatedBuildInputs = [ (lib.getBin mpi) ];
+  propagatedBuildInputs = [
+    (lib.getBin mpi)
+  ] ++ lib.optionals enablePython (with pythonPackages; [
+    mpi4py
+    voluptuous
+    numpy
+    h5py
+    scipy
+    pyyaml
+  ]);
 
   CXXFLAGS = [
     # GCC 13: error: 'uintptr_t' in namespace 'std' does not name a type
@@ -92,20 +108,20 @@ stdenv.mkDerivation rec {
   ];
 
   cmakeFlags = [
-    "-DUSE_SCALAPACK=ON"
+    "-DSIRIUS_USE_SCALAPACK=ON"
+    "-DSIRIUS_USE_VDWXC=ON"
+    "-DSIRIUS_CREATE_FORTRAN_BINDINGS=ON"
+    "-DSIRIUS_USE_OPENMP=ON"
     "-DBUILD_TESTING=ON"
-    "-DUSE_VDWXC=ON"
-    "-DCREATE_FORTRAN_BINDINGS=ON"
-    "-DUSE_OPENMP=ON"
-    "-DBUILD_TESTING=ON"
-  ]
-  ++ lib.optionals (gpuBackend == "cuda") [
-    "-DUSE_CUDA=ON"
+  ] ++ lib.optionals (gpuBackend == "cuda") [
+    "-DSIRIUS_USE_CUDA=ON"
     "-DCUDA_TOOLKIT_ROOT_DIR=${cudaPackages.cudatoolkit}"
-  ]
-  ++ lib.optionals (gpuBackend == "rocm") [
-    "-DUSE_ROCM=ON"
+    (lib.cmakeFeature "CMAKE_CUDA_ARCHITECTURES" cudaPackages.flags.cmakeCudaArchitecturesString)
+  ] ++ lib.optionals (gpuBackend == "rocm") [
+    "-DSIRIUS_USE_ROCM=ON"
     "-DHIP_ROOT_DIR=${rocmPackages.clr}"
+  ] ++ lib.optionals enablePython [
+    "-DSIRIUS_CREATE_PYTHON_MODULE=ON"
   ];
 
   doCheck = true;

pkgs/by-name/um/umpire/package.nix

@@ -2,8 +2,13 @@
 , lib
 , fetchFromGitHub
 , cmake
+, config
+, cudaSupport ? config.cudaSupport
+, cudaPackages ? null
 }:
 
+assert cudaSupport -> cudaPackages != null;
+
 stdenv.mkDerivation rec {
   pname = "umpire";
   version = "2024.07.0";
@@ -16,7 +21,22 @@ stdenv.mkDerivation rec {
     fetchSubmodules = true;
   };
 
-  nativeBuildInputs = [ cmake ];
+  nativeBuildInputs = [
+    cmake
+  ] ++ lib.optionals cudaSupport [
+    cudaPackages.cuda_nvcc
+  ];
+
+  buildInputs = lib.optionals cudaSupport (with cudaPackages; [
+    cudatoolkit
+    cuda_cudart
+  ]);
+
+  cmakeFlags = lib.optionals cudaSupport [
+    "-DCUDA_TOOLKIT_ROOT_DIR=${cudaPackages.cudatoolkit}"
+    "-DENABLE_CUDA=ON"
+    (lib.cmakeFeature "CMAKE_CUDA_ARCHITECTURES" cudaPackages.flags.cmakeCudaArchitecturesString)
+  ];
 
   meta = with lib; {
     description = "Application-focused API for memory management on NUMA & GPU architectures";

pkgs/top-level/python-packages.nix

@@ -14286,6 +14286,11 @@ self: super: with self; {
 
   sipyco = callPackage ../development/python-modules/sipyco { };
 
+  sirius = toPythonModule (pkgs.sirius.override {
+    enablePython = true;
+    pythonPackages = self;
+  });
+
   sismic = callPackage ../development/python-modules/sismic { };
 
   sisyphus-control = callPackage ../development/python-modules/sisyphus-control { };