Fix OpenMP build failure on macOS

environment.yaml

@@ -8,3 +8,4 @@ dependencies:
   - numpy<=2.0
   - horton
   - h5py
+  - llvm-openmp

setup.py

@@ -1,13 +1,35 @@
+import os
+import platform
+
 import pybind11
 from setuptools import Extension, setup
 
+
+def get_openmp_flags():
+    if platform.system() == "Darwin":
+        # Apple clang does not support -fopenmp directly.
+        conda_prefix = os.environ.get("CONDA_PREFIX", "")
+        if not conda_prefix:
+            raise RuntimeError(
+                "CONDA_PREFIX is not set. Activate your conda environment first."
+            )
+        compile_args = ["-Xpreprocessor", "-fopenmp", f"-I{conda_prefix}/include"]
+        link_args = [f"-L{conda_prefix}/lib", "-lomp"]
+    else:
+        compile_args = ["-fopenmp"]
+        link_args = ["-fopenmp"]
+    return compile_args, link_args
+
+
+_omp_compile, _omp_link = get_openmp_flags()
+
 ext_modules = [
     Extension(
         "pyxdm.core.exchange_hole_cpp",
         ["pyxdm/cpp/exchange_hole.cpp"],
         include_dirs=[pybind11.get_include()],
-        extra_compile_args=["-O3", "-fopenmp"],
-        extra_link_args=["-fopenmp"],
+        extra_compile_args=["-O3"] + _omp_compile,
+        extra_link_args=_omp_link,
         language="c++",
     )
 ]