Fix TEM build using f2py's meson backend

external_modules/total_energy_module/build_total_energy_module.sh

@@ -1,12 +1,6 @@
 #!/bin/bash
 
-# usage:
-#
-#   $ ./this.sh
-# or
-#   $ F2PY=/usr/bin/f2py3 ./this.sh
-
-set -u
+set -ue
 
 err(){
     echo "error $@"
@@ -15,6 +9,23 @@ err(){
 
 [ $# -eq 1 ] || err "Please provide exactly one argument (the path to the QE directory)" && root_dir=$1
 
+
+# ----------------------------------------------------------------------------
+# Settings to adapt to your machine
+# ----------------------------------------------------------------------------
+
+# default: system blas,lapack and fftw, adapt as needed
+linalg="-lblas -llapack"
+fftw="-lfftw3"
+
+# This is for OpenMPI. Intel MPI's compiler wrappers are called mpiifort and
+# mpiicc. Both libraries provide other aliases like mpif90 and mpicc.
+f_compiler="mpifort"
+c_compiler="mpicc"
+
+# ----------------------------------------------------------------------------
+
+
 echo "Using QE root dir: $root_dir"
 
 pw_src_path=$root_dir/PW/src
@@ -32,45 +43,53 @@ modobjs="$(find $root_dir/Modules/ -name "*.o")"
 [ -n "$utilobjs" ] || err "utilobjs empty"
 [ -n "$modobjs" ] || err "modobjs empty"
 
-project_lib_folders=" -L$root_dir/Modules -L$root_dir/KS_Solvers -L$root_dir/FFTXlib/src -L$root_dir/LAXlib -L$root_dir/UtilXlib -L$root_dir/dft-d3 -L$root_dir/upflib -L$root_dir/XClib -L$root_dir/external/devxlib/src -L$root_dir/external/mbd/src"
-project_libs="-lqemod -lks_solvers -lqefft -lqela -lutil -ldftd3qe -lupf -ldevXlib -lmbd"
-project_inc_folders="-I$root_dir/Modules -I$root_dir/FFTXlib/src -I$root_dir/LAXlib -I$root_dir/KS_Solvers -I$root_dir/UtilXlib -I$root_dir/upflib -I$root_dir/XClib -I$root_dir/external/devxlib/src -I$root_dir/external/mbd/src"
 
-# default: system blas,lapack and fftw, adapt as needed
-linalg="-lblas -llapack"
-fftw="-lfftw3"
+# -lfoo will link shared libs (libfoo.so) as well as static ones (libfoo.a). In
+# qe_libs we list the ones we need in total_energy (which are almost all of
+# them anyway), but since we are lazy, we use all paths to them as link (-L)
+# and include (-I) dirs.
+qe_static_lib_dirs=$(find $root_dir -name "*.a" | xargs dirname | sort -u | paste -s -d '@')
+qe_lib_dirs=-L$(echo "$qe_static_lib_dirs" | sed -re 's/@/ -L/g')
+qe_inc_dirs=-I$(echo "$qe_static_lib_dirs" | sed -re 's/@/ -I/g')
+qe_libs="-lqemod -lks_solvers -lqefft -lqela -lutil -ldftd3qe -lupf -ldevXlib -lmbd"
 
+echo "qe_lib_dirs: $qe_lib_dirs"
+echo "qe_inc_dirs: $qe_inc_dirs"
 
 here=$(pwd)
 mod_name=total_energy
-src=${mod_name}.f90
-#cp -v $src $pw_src_path/
-#cd $pw_src_path
-rm -vf ${mod_name}.*.so
-
-# I think the -I$pw_src_path is necessary now or else
-# the compiler won't find the .mod files anymore since
-# the build is done in some temporary directory.
-# As far as I understand it, the issue lies with the compiler
-# ignoring the *.o files.
-FC="mpif90" python3 -m numpy.f2py \
-    -c $src \
+tmp_lib_dir=$here/libs
+meson_builddir=$here/meson_builddir
+rm -rf $tmp_lib_dir $meson_builddir
+
+mkdir -p $tmp_lib_dir
+# xc_lib.a needs to be called lib<someting>.a to be linkable with -l<something>
+cp $root_dir/XClib/xc_lib.a $here/libs/libxclib.a
+
+# Stick all object files into a static lib so that we can link them, just
+# specifying them on the f2py CLI doesn't work with the meson backend, as it
+# seems.
+ar rcs $tmp_lib_dir/liballobjs.a $pwobjs $utilobjs $modobjs
+
+FC="$f_compiler" \
+CC="$c_compiler" \
+FFLAGS="-I$tmp_lib_dir -I$pw_src_path $qe_inc_dirs" \
+LDFLAGS="-L$tmp_lib_dir -L$pw_src_path $qe_lib_dirs" \
+python3 -m numpy.f2py \
+    --backend meson \
+    --dep mpi \
+    --build-dir $meson_builddir \
+    -c ${mod_name}.f90 \
     -m $mod_name \
-    $project_inc_folders \
-    $project_lib_folders \
-    $project_libs \
-    $pwobjs \
-    $utilobjs \
-    $modobjs \
     $linalg \
     $fftw \
-    $root_dir/XClib/xc_lib.a \
-    -I$pw_src_path \
-
-#${F2PY:=f2py} \
-#    --f90exec=mpif90 \
-#    --f77exec=mpif90 \
-
-#rm $src
-#mv -v ${mod_name}.*.so $here/
-#cd $here
+    $qe_libs \
+    -lallobjs -lxclib
+
+# Workaround for f2py+meson+mpi bug
+# (https://github.com/numpy/numpy/issues/28902)
+cd $meson_builddir
+sed -i -re "s/dependency\('mpi'\)/dependency\('mpi', language: 'fortran')/" meson.build
+meson setup --reconfigure bbdir
+meson compile -C bbdir
+cp -v bbdir/${mod_name}.*.so $here/

requirements.txt

@@ -3,9 +3,14 @@
 # such that a pre-installed torch is not found.
 torch
 
+# We need a recent numpy to build external_modules/total_energy_module using
+# numpy's f2py. See
+# https://github.com/mala-project/mala/issues/559#issuecomment-2864133103
+# for details.
+numpy >= 2
+
 ase
 mpmath
-numpy
 optuna
 scipy
 pandas