v0.1.6 upload

- updated scripts for the MSORF preprint v2: https://arxiv.org/abs/2505.21247v2

- added tests for MSORF

- revised SC Huber implementation

- added an environmental variable for avoiding Numba-Numpy parallelization (see https://stackoverflow.com/questions/79673925/disable-numpy-parallelization-inside-numba-jit)

- minor code revision

Author: kvkarandashev

.pre-commit-config.yaml

@@ -77,3 +77,8 @@ repos:
         -   id: conventional-pre-commit
             stages: [commit-msg]
             args: []
+
+    -   repo: https://github.com/kieran-ryan/pyprojectsort
+        rev: v0.4.0
+        hooks:
+        -   id: pyprojectsort

README.md

@@ -16,11 +16,13 @@ or, if `makefile` is installed,
    ```
 Some parts of the code depend on additional dependencies that can be installed with the corresponding optional dependecies flag. The three defined for this package are
 
-- `orb_ml` - for FJK (machine learning from orbital information)
+- `orb_ml` - for FJK (machine learning from orbital information).
 
-- `msorf` - for MSORF (everything in `qml2.multilevel_sorf`)
+- `msorf` - for MSORF (everything in `qml2.multilevel_sorf`).
 
-- `torch` - Torch functionality (efficiency questionable right now TBH)
+- `morfeus` - for applications dependend on `morfeus-ml` package (everything related to conformer ensemble generation).
+
+- `torch` - Torch functionality (efficiency questionable right now TBH).
 
 For example, to use the `orb_ml` and `msorf` optional dependecy flags in your installation use
 
@@ -42,6 +44,8 @@ To check that the installed repo works correctly run
    ```
 **NOTE:** The command assumes that `python` environmental variable points towards a valid executable. If you use an environment alias change definition of the `python` variable in the beginning of the Makefile.
 
+If you have trouble getting all tests to complete consult the `tests/qml2_test.yml` file describing the conda environment where all tests have been created and reproduced.
+
 ## :books: API documentation
 
 API documentation can be generated with [Doxygen](https://www.doxygen.nl/) by running
@@ -64,6 +68,10 @@ This will create `manual.html` file that can be opened with an Internet browser.
 
 `QML2_NUM_PROCS` - number of processes spawned by parts of the code parallelized via `python.multiprocessing` (training set representations in model-related classes, `pyscf` calculations made by `OML_Compound_list` attribute calls). For limiting number of OpenMP threads spawned in turn by these processes use suppress options (such as `KRRModel` class's `training_reps_suppress_openmp` option). Also see `parallelization.set_default_num_procs`.
 
+`QML2_NUM_PROCS` - number of processes spawned by parts of the code parallelized via `python.multiprocessing` (training set representations in model-related classes, `pyscf` calculations made by `OML_Compound_list` attribute calls). For limiting number of OpenMP threads spawned in turn by these processes use suppress options (such as `KRRModel` class's `training_reps_suppress_openmp` option). Also see `parallelization.set_default_num_procs`.
+
+`QML2_AVOID_NUMBA_NUMPY_PARALLELIZATION` - some Numba routines in the code call in parallel Numpy routines, which creates problems in some setups (e.g. when both Numba and Numpy try to parallelize over a large number of threads without taking each other into account). Setting this environmental variable to `1` disables Numba parallelization in such routines, leaving them to be parallelized exclusively with Numpy.
+
 ### Experimental
 
 `QML2_DEFAULT_JIT` - setting to `NUMBA` (default) or `TORCH` (both are case insensitive) determines whether Numba or TorchScript JIT compilation is used. Also see `jit_interfaces.set_default_jit`.

misc_commands/spython

@@ -8,6 +8,8 @@
 # 1st argument - python script name.
 # 2nd argument - how the SLURM job is named.
 # Arguments starting from third (if present) - arguments for the python script.
+# If SPYTHON_NO_SLURM is defined in shell as 1 "spython" launches the submitted script in simple shell. Introduced in order same bash scripts could be used to spam jobs on workstations that do not have SLURM.
+
 
 #TODO : K.Karan: I think GPU allocation is OK according to the manuals, but I never got it to work
 # on my workstation.
@@ -16,7 +18,19 @@ function flag_val(){
     echo $1 | cut -d'=' -f2-
 }
 
-NCPUs=$(cat /etc/slurm/slurm.conf | tr ' ' '\n' | awk 'BEGIN {FS="="} {if ($1=="CPUs") print $2}')
+slurmconf_file="/etc/slurm/slurm.conf"
+
+if [ -f "$slurmconf_file" ]
+then
+    NCPUs=$(cat $slurmconf_file | tr ' ' '\n' | awk 'BEGIN {FS="="} {if ($1=="CPUs") print $2}')
+else
+    echo "WARNING: SLURM configuration file not found."
+    NCPUs=$OMP_NUM_THREADS
+    if [ "$NCPUs" == "" ]
+    then
+        NCPUs=1
+    fi
+fi
 NGPUs=0
 JOB_OMP_NUM_THREADS=$NCPUs
 nodelist=$(hostname)
@@ -26,6 +40,11 @@ conda_env=$CONDA_DEFAULT_ENV
 PIPE_ARGS=0
 array=""
 
+if [ "$SPYTHON_NO_SLURM" == "" ]
+then
+    SPYTHON_NO_SLURM=0
+fi
+
 while [[ "$1" == --* ]]
 do
     case "$1" in
@@ -196,9 +215,14 @@ then
 fi
 EOF
 
-SLURM_ID=$(sbatch $array_arg $DEPSTRING $EXNAME | awk '{print $4}')
-
-DEPSTRING="--dependency=afterok:$SLURM_ID"
+if [ "$SPYTHON_NO_SLURM" == "0" ]
+then
+    SLURM_ID=$(sbatch $array_arg $DEPSTRING $EXNAME | awk '{print $4}')
+    DEPSTRING="--dependency=afterok:$SLURM_ID"
+else
+    chmod +x $EXNAME
+    ./$EXNAME > $job_name.stdout_0 2> $job_name.stderr_0
+fi
 
 cd ..
 

paper_scripts/2025_multilevel_sorf.tar.gz

@@ -1,41 +1,63 @@
 [build-system]
-requires = ["setuptools >= 61.0"]
 build-backend = "setuptools.build_meta"
+requires = [
+    "setuptools >= 61.0",
+]
+
 [project]
+authors = [
+    { name = "Danish Khan" },
+    { name = "Jan Weinreich" },
+    { name = "Konstantin Karandashev" },
+    { name = "Stefan Heinen" },
+]
+dependencies = [
+    "numba",
+    "numpy",
+    "scipy",
+]
+description = "Collection of Quantum Machine Learning routines"
 name = "qml2"
-version = "0.1.5"
-dependencies=[
-  "numpy",
-  "numba",
-  "scipy",
-]
-authors=[
- {name="Konstantin Karandashev"},
- {name="Stefan Heinen"},
- {name="Danish Khan"},
- {name="Jan Weinreich"},
-]
-description="Collection of Quantum Machine Learning routines"
-readme="README.md"
-license={text="MIT License"}
-[tool.setuptools]
-packages=[
-  "qml2",
-  "qml2.orb_ml",
-  "qml2.representations",
-  "qml2.models",
-  "qml2.jit_interfaces",
-  "qml2.dataset_formats",
-  "qml2.kernels",
-  "qml2.test_utils",
-  "qml2.gpytorch",
-  "qml2.cupy",
-  "qml2.multilevel_sorf",
-  "qml2.multilevel_sorf.processed_object_constructors",
-  "qml2.multilevel_sorf.base_constructors"]
+readme = "README.md"
+version = "0.1.6"
+
+[project.license]
+text = "MIT License"
+
 [project.optional-dependencies]
-torch=["torch"]
-orb_ml=["pyscf", "geometric"]
-msorf=["aalto-boss"]
+morfeus = [
+    "morfeus-ml",
+    "rdkit",
+]
+msorf = [
+    "aalto-boss",
+]
+orb_ml = [
+    "geometric",
+    "pyscf",
+]
+torch = [
+    "torch",
+]
+
+[tool.setuptools]
+packages = [
+    "qml2",
+    "qml2.cupy",
+    "qml2.dataset_formats",
+    "qml2.gpytorch",
+    "qml2.jit_interfaces",
+    "qml2.kernels",
+    "qml2.models",
+    "qml2.multilevel_sorf",
+    "qml2.multilevel_sorf.base_constructors",
+    "qml2.multilevel_sorf.processed_object_constructors",
+    "qml2.orb_ml",
+    "qml2.representations",
+    "qml2.test_utils",
+]
+
 [tool.setuptools.package-data]
-qml2 = ["*/*.cu"]
+qml2 = [
+    "*/*.cu",
+]

pyproject.toml

@@ -3,6 +3,7 @@
 import os
 from copy import deepcopy
 
+# TODO make morfeus and rdkit dependences here optional to allow running some MSORF tests without morfeus-ml
 from morfeus.conformer import ConformerEnsemble
 from rdkit import Chem, RDLogger
 
@@ -14,6 +15,7 @@
 from .orb_ml.representations import OML_rep_params
 from .utils import weighted_array, write_compound_to_xyz_file
 
+# disable RdKit's verbose mode.
 RDLogger.DisableLog("rdApp.*")
 
 base_compound_class_dict = {
@@ -23,6 +25,26 @@
     "OML_Slater_pairs": OML_Slater_pairs,
 }
 
+morfeus_random_seed = None
+
+
+def set_morfeus_random_seed(new_morfeus_random_seed):
+    """
+    Set RNG seed used by morfeus when generating conformers for the Ensemble class.
+    """
+    global morfeus_random_seed
+    morfeus_random_seed = new_morfeus_random_seed
+
+
+def get_morfeus_random_seed(default_seed=None):
+    """
+    Get RNG seed used by morfeus when generating conformers for the Ensemble class.
+    """
+    if default_seed is None:
+        return morfeus_random_seed
+    else:
+        return default_seed
+
 
 class WeightedCompound:
     def __init__(self, compound, rho, ff_energy):
@@ -53,6 +75,7 @@ def __init__(
         num_conformer_generations=1,
         savefile_prefix=None,
         compound_kwargs={},
+        random_seed=None,
     ):
         self.rdkit_obj = rdkit_obj
         self.SMILES = SMILES
@@ -70,6 +93,7 @@ def __init__(
         self.filtered_conformers = None
         self.processed_conformers = None
         self.savefile_prefix = savefile_prefix
+        self.random_seed = random_seed
 
     def get_nuclear_charges(self):
         return array_([a.GetAtomicNum() for a in self.rdkit_obj.GetAtoms()])
@@ -113,6 +137,7 @@ def conformer_generation(self):
                 n_confs=self.num_conformers,
                 optimize=self.ff_type,
                 n_threads=checked_environ_val("MORFEUS_NUM_THREADS", default_answer=1),
+                random_seed=get_morfeus_random_seed(default_seed=self.random_seed),
             )
         except Exception as ex:
             if not isinstance(ex, ValueError):

qml2/ensemble.py

@@ -12,10 +12,11 @@
 # then it enters under the name given to it by Numpy/Numba.
 import importlib
 
-from ..basic_utils import checked_environ_val
+from ..basic_utils import checked_environ_val, checked_logical_environ_val
 from .jit_manager import available_jits, numba_flag
 
 default_jit_env_var_name = "QML2_DEFAULT_JIT"
+avoid_numba_numpy_parallellization_env_var_name = "QML2_AVOID_NUMBA_NUMPY_PARALLELIZATION"
 
 
 def set_defaults_from_interface(interface_name):
@@ -36,6 +37,15 @@ def set_default_jit(new_jit_flag):
     set_defaults_from_interface("." + used_jit_name + "_interface")
 
 
+def allow_numba_numpy_parallelization():
+    """
+    This function is used as the `numba_parallel` value in `jit_` instances for routines where `prange` loop includes references to Numpy.
+    """
+    return not checked_logical_environ_val(
+        avoid_numba_numpy_parallellization_env_var_name, default_answer=False
+    )
+
+
 if __name__ != "__main__":
     # Check environment for default jit flag, if none pick numba
     default_flag = checked_environ_val(

qml2/jit_interfaces/__init__.py

@@ -1,5 +1,6 @@
 # Introduced to ensure flags for different JIT compilers do not obstruct each other.
 import inspect
+import traceback
 
 from ..basic_utils import checked_logical_environ_val
 
@@ -46,7 +47,7 @@ class JITFailureException(Exception):
 
 
 class ExceptionRaisingFunc:
-    def __init__(self, func):
+    def __init__(self, func, ex):
         self.exception_text = (
             """
 Failure in JIT compilation.
@@ -56,6 +57,10 @@ def __init__(self, func):
 function source:
 """
             + inspect.getsource(func)
+            + """
+traceback:
+"""
+            + "\n".join(traceback.format_exception(ex))
         )
 
     def __call__(self, *args, **kwargs):
@@ -96,10 +101,10 @@ def exception_skipping_jit(self, signature_or_function, local_skip_jit=False, **
             return signature_or_function
         try:
             return self.used_jit_func(signature_or_function, **kwargs)
-        except (*self.possible_jit_failures,):
+        except (*self.possible_jit_failures,) as ex:
             #            assert type(ex) in self.possible_jit_failures
             if skip_jit_failures:
-                return ExceptionRaisingFunc(signature_or_function)
+                return ExceptionRaisingFunc(signature_or_function, ex)
             else:
                 raise JITFailureException
 

qml2/jit_interfaces/jit_manager.py

@@ -125,8 +125,43 @@ def random_array_from_rng_(size=(1,), rng=None):
         return rng.random(size)
 
 
+@jit_
+def random_from_rng_(rng=None):
+    return random_array_from_rng_(rng=rng)[0]
+
+
 standard_normal_ = np.random.standard_normal
+
+
+@jit_
+def standard_normal_array_from_rng_(size=(1,), rng=None):
+    if rng is None:
+        return standard_normal_(size)
+    else:
+        return rng.standard_normal(size)
+
+
+@jit_
+def standard_normal_from_rng_(rng=None):
+    return standard_normal_array_from_rng_(rng=rng)[0]
+
+
 randint_ = np.random.randint
+
+
+@jit_
+def randint_array_from_rng_(lbound, ubound, size=(1,), rng=None):
+    if rng is None:
+        return randint_(lbound, ubound, size=size)
+    else:
+        return rng.integers(lbound, ubound, size=size)
+
+
+@jit_
+def randint_from_rng_(lbound, ubound, rng=None):
+    return randint_array_from_rng_(lbound, ubound, rng=rng)[0]
+
+
 seed_ = np.random.seed
 permutation_ = np.random.permutation
 # copying