Upload 686 files

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+env/SE3Transformer/images/se3-transformer.png filter=lfs diff=lfs merge=lfs -text
+img/diffusion_protein_gradient_2.jpg filter=lfs diff=lfs merge=lfs -text
+pyrosetta-2023.14+release.7132bdc754a-cp310-cp310-linux_x86_64.whl filter=lfs diff=lfs merge=lfs -text

.github/CODEOWNERS

@@ -0,0 +1,2 @@
+# Benchmark scripts
+/.rosetta-ci @lyskov

.gitignore

@@ -0,0 +1,16 @@
+*.py[cod]
+rfdiffusion.egg-info
+
+models/
+schedules/
+
+examples/ppi_scaffolds
+
+tests/.results.json
+tests/input_pdbs
+tests/outputs
+tests/ppi_scaffolds
+tests/reference_outputs/
+tests/target_folds
+tests/tim_barrel_scaffold
+tests/tests_*

.rosetta-ci/.gitignore

@@ -0,0 +1,3 @@
+*.pyc
+results/
+benchmark.ubuntu.ini

.rosetta-ci/benchmark.py

@@ -0,0 +1,410 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+# (c) Copyright Rosetta Commons Member Institutions.
+# (c) This file is part of the Rosetta software suite and is made available under license.
+# (c) The Rosetta software is developed by the contributing members of the Rosetta Commons.
+# (c) For more information, see http://www.rosettacommons.org. Questions about this can be
+# (c) addressed to University of Washington CoMotion, email: license@uw.edu.
+
+## @file   benchmark.py
+## @brief  Run arbitrary Rosetta testing script
+## @author Sergey Lyskov
+
+from __future__ import print_function
+
+import os, os.path, sys, shutil, json, platform, re
+import codecs
+
+from importlib.machinery import SourceFileLoader
+
+from configparser import ConfigParser, ExtendedInterpolation
+import argparse
+
+from tests import *  # execute, Tests states and key names
+from hpc_drivers import *
+
+
+# Calculating value of Platform dict
+Platform = {}
+if sys.platform.startswith("linux"):
+    Platform['os'] = 'ubuntu' if os.path.isfile('/etc/lsb-release') and 'Ubuntu' in open('/etc/lsb-release').read() else 'linux'  # can be linux1, linux2, etc
+elif sys.platform == "darwin" :      Platform['os'] = 'mac'
+elif sys.platform == "cygwin" :      Platform['os'] = 'cygwin'
+elif sys.platform == "win32" :       Platform['os'] = 'windows'
+else:                                Platform['os'] = 'unknown'
+
+#Platform['arch'] = platform.architecture()[0][:2]  # PlatformBits
+Platform['compiler'] = 'gcc' if Platform['os'] == 'linux' else 'clang'
+
+Platform['python'] = sys.executable
+
+
+def load_python_source_from_file(module_name, module_path):
+    ''' replacment for deprecated imp.load_source
+    '''
+    return SourceFileLoader(module_name, module_path).load_module()
+
+
+class Setup(object):
+    __slots__ = 'test working_dir platform config compare debug'.split()  # version daemon path_to_previous_test
+    def __init__(self, **attrs):
+        #self.daemon = True
+        for k, v in attrs.items():
+            if k in self.__slots__: setattr(self, k, v)
+
+
+def setup_from_options(options):
+    ''' Create Setup object based on user supplied options, config files and auto-detection
+    '''
+    platform = dict(Platform)
+
+    if options.suffix: options.suffix = '.' + options.suffix
+
+    platform['extras'] = options.extras.split(',') if options.extras else []
+    platform['python'] = options.python
+    #platform['options'] = json.loads( options.options ) if options.options else {}
+
+    if options.memory: memory = options.memory
+    elif platform['os'] in ['linux', 'ubuntu']: memory = int( execute('Getting memory info...', 'free -m', terminate_on_failure=False, silent=True, silence_output_on_errors=True, return_='output').split('\n')[1].split()[1]) // 1024
+    elif platform['os'] == 'mac':   memory = int( execute('Getting memory info...', 'sysctl -a | grep hw.memsize', terminate_on_failure=False, silent=True, silence_output_on_errors=True, return_='output').split()[1]) // 1024 // 1024 // 1024
+
+    platform['compiler'] = options.compiler
+
+    if os.path.isfile(options.config):
+        with open(options.config) as f:
+            if '%(here)s' in f.read():
+                print(f"\n\n>>> ERROR file `{options.config}` seems to be in outdated format! Please use benchmark.template.ini to update it.")
+                sys.exit(1)
+
+        user_config = ConfigParser(
+            dict(
+                _here_ = os.path.abspath('./'),
+                _user_home_ = os.environ['HOME']
+            ),
+            interpolation = ExtendedInterpolation()
+        )
+
+        with open(options.config) as f: user_config.readfp(f)
+
+    else:
+        print(f"\n\n>>> Config file `{options.config}` not found. You may want to manually copy `benchmark.ini.template` to `{options.config}` and edit the settings\n\n")
+        user_config = ConfigParser()
+        user_config.set('main', 'cpu_count',  '1')
+        user_config.set('main', 'hpc_driver', 'MultiCore')
+        user_config.set('main', 'branch',     'unknown')
+        user_config.set('main', 'revision',   '42')
+        user_config.set('main', 'user_name',  'Jane Roe')
+        user_config.set('main', 'user_email', 'jane.roe@university.edu')
+        user_config.add_section('main')
+
+    if options.jobs: user_config.set('main', 'cpu_count', str(options.jobs) )
+    user_config.set('main', 'memory',    str(memory) )
+
+    if options.mount:
+        for m in options.mount:
+            key, _, path = m.partition(':')
+            user_config.set('mount', key, path)
+
+    #config = Config.items('config')
+    #for section in config.sections(): print('Config section: ', section, dict(config.items(section)))
+    #config = { section: dict(Config.items(section)) for section in Config.sections() }
+
+    config = { k : d for k, d in user_config['main'].items() if k not in user_config[user_config.default_section] }
+    config['mounts'] = { k : d for k, d in user_config['mount'].items() if k not in user_config[user_config.default_section] }
+
+    #print(json.dumps(config, sort_keys=True, indent=2)); sys.exit(1)
+
+    #config.update( config.pop('config').items() )
+
+    config = dict(config,
+                  cpu_count = user_config.getint('main', 'cpu_count'),
+                  memory = memory,
+                  revision = user_config.getint('main', 'revision'),
+                  emulation=True,
+    )  # debug=options.debug,
+
+    if 'results_root' not in config: config['results_root'] = os.path.abspath('./results/')
+
+    if 'prefix' in config:
+        assert os.path.isabs( config['prefix'] ), f'ERROR: `prefix` path must be absolute! Got: {config["prefix"]}'
+
+    else: config['prefix'] = os.path.abspath( config['results_root'] + '/prefix')
+
+    config['merge_head'] = options.merge_head
+    config['merge_base'] = options.merge_base
+
+    if options.skip_compile is not None: config['skip_compile'] = options.skip_compile
+
+    #print(f'Results path: {config["results_root"]}')
+    #print('Config:{}, Platform:{}'.format(json.dumps(config, sort_keys=True, indent=2), Platform))
+
+    if options.compare: print('Comparing tests {} with suffixes: {}'.format(options.args, options.compare) )
+    else: print('Running tests: {}'.format(options.args) )
+
+    if len(options.args) != 1: print('Error: Single test-name-to-run should be supplied!');  sys.exit(1)
+    else:
+        test = options.args[0]
+        if test.startswith('tests/'): test = test.partition('tests/')[2][:-3]  # removing dir prefix and .py suffix
+
+    if options.compare:
+        compare = options.compare[0], options.compare[1]  # (this test suffix, previous test suffix)
+        working_dir = os.path.abspath( config['results_root'] + f'/{platform["os"]}.{test}' )  # will be a root dir with sub-dirs (options.compare[0], options.compare[1])
+    else:
+        compare = None
+        working_dir = os.path.abspath( config['results_root'] + f'/{platform["os"]}.{test}{options.suffix}' )
+
+
+    if os.path.isdir(working_dir): shutil.rmtree(working_dir);  #print('Removing old job dir %s...' % working_dir)  # remove old dir if any
+    os.makedirs(working_dir)
+
+    setup = Setup(
+        test        = test,
+        working_dir = working_dir,
+        platform    = platform,
+        config      = config,
+        compare     = compare,
+        debug       = options.debug,
+        #daemon      = False,
+    )
+
+    setup_as_json = json.dumps( { k : getattr(setup, k) for k in setup.__slots__}, sort_keys=True, indent=2)
+    with open(working_dir + '/.setup.json', 'w') as f: f.write(setup_as_json)
+
+    #print(f'Detected hardware platform: {Platform}')
+    print(f'Setup: {setup_as_json}')
+    return setup
+
+
+def truncate_log(log):
+    _max_log_size_  = 1024*1024*1
+    _max_line_size_ = _max_log_size_ // 2
+
+    if len(log) > _max_log_size_:
+        new = log
+        lines = log.split('\n')
+
+        if len(lines) > 256:
+            new_lines = lines[:32] + ['...truncated...'] + lines[-128:]
+            new = '\n'.join(new_lines)
+
+        if len(new) > _max_log_size_: # special case for Ninja logs that does not use \n
+            lines = re.split(r'[\r\n]*', log)  #t.log.split('\r')
+            if len(lines) > 256: new = '\n'.join( lines[:32] + ['...truncated...'] + lines[-128:] )
+
+        if len(new) > _max_log_size_: # going to try to truncate each individual line...
+            print(f'Trying to truncate log line-by-line...')
+            new = '\n'.join( (
+                ( line[:_max_line_size_//3] + '...truncated...' + line[-_max_line_size_//3:] ) if line > _max_line_size_ else line
+                for line in new_lines ) )
+
+        if len(new) > _max_log_size_: # fall-back strategy in case all of the above failed...
+            print(f'WARNING: could not truncate log line-by-line, falling back to raw truncate...')
+            new = 'WARNING: could not truncate test log line-by-line, falling back to raw truncate!\n...truncated...\n' + ( '\n'.join(lines) )[-_max_log_size_+256:]
+
+        print( 'Trunacting test output log: {0}MiB --> {1}MiB'.format(len(log)/1024/1024, len(new)/1024/1024) )
+
+        log = new
+
+    return log
+
+def truncate_results_logs(results):
+    results[_LogKey_] = truncate_log( results[_LogKey_] )
+    if _ResultsKey_ in results  and  _TestsKey_ in results[_ResultsKey_]:
+        tests = results[_ResultsKey_][_TestsKey_]
+        for test in tests:
+            tests[test][_LogKey_] = truncate_log( tests[test][_LogKey_] )
+
+
+def find_test_description(test_name, test_script_file_name):
+    ''' return content of test-description file if any or None if no description was found
+    '''
+
+    def find_description_file(prefix, test_name):
+        fname = prefix + test_name + '.md'
+        if os.path.isfile(fname): return fname
+        return prefix + 'md'
+
+    description_file_name =  find_description_file( test_script_file_name[:-len('command.py')] + 'description.', test_name) if test_script_file_name.endswith('/command.py') else find_description_file(test_script_file_name[:-len('py')], test_name)
+
+    if description_file_name  and  os.path.isfile(description_file_name):
+        print(f'Found test suite description in file: {description_file_name!r}')
+        with open(description_file_name, encoding='utf-8', errors='backslashreplace') as f: description = f.read()
+        return description
+
+    else: return None
+
+
+
+def run_test(setup):
+    #print(f'{setup!r}')
+    suite, rest = setup.test.split('.'), []
+    while suite:
+        #print( f'suite: {suite}, test: {rest}' )
+
+        file_name = '/'.join( ['tests'] + suite ) + '.py'
+        if os.path.isfile(file_name): break
+
+        file_name = '/'.join( ['tests'] + suite ) + '/command.py'
+        if os.path.isfile(file_name): break
+
+        rest.insert(0, suite.pop())
+
+
+    test = '.'.join( suite + rest )
+    test_name = '.'.join(rest)
+
+    print( f'Loading test from: {file_name}, suite+test: {test!r}, test: {test_name!r}' )
+    #test_suite = imp.load_source('test_suite', file_name)
+    test_suite = load_python_source_from_file('test_suite', file_name)
+
+    test_description = find_test_description(test_name, file_name)
+
+    if setup.compare:
+        #working_dir_1 = os.path.abspath( config['results_root'] + f'/{Platform["os"]}.{test}.{Options.compare[0]}' )
+        working_dir_1 = setup.working_dir + f'/{setup.compare[0]}'
+
+        working_dir_2        = setup.compare[1]  and  ( setup.working_dir + f'/{setup.compare[1]}' )
+        res_2_json_file_path = setup.compare[1]  and  f'{working_dir_2}/.execution.results.json'
+
+        with open(working_dir_1 + '/.execution.results.json') as f: res_1 = json.load(f).get(_ResultsKey_)
+
+        if setup.compare[1] and ( not os.path.isfile(res_2_json_file_path) ):
+            setup.compare[1] = None
+            state_override = _S_failed_
+        else:
+            state_override = None
+
+        if setup.compare[1] == None: res_2, working_dir_2 = None, None
+        else:
+            with open(res_2_json_file_path) as f: res_2 = json.load(f).get(_ResultsKey_)
+
+        res = test_suite.compare(test, res_1, working_dir_1, res_2, working_dir_2)
+
+        if state_override:
+            log_prefix = \
+                f'WARNING: Previous test results does not have `.execution.results.json` file, so comparision with None was performed instead!\n' \
+                f'WARNING: Overriding calcualted test state `{res[_StateKey_]}` → `{_S_failed_}`...\n\n'
+
+            res[_LogKey_] = log_prefix + res[_LogKey_]
+            res[_StateKey_] = _S_failed_
+
+
+        # # Caution! Some of the strings in the result object may be unicode. Be robust to unicode in the log messages.
+        # with codecs.open(setup.working_dir+'/.comparison.log.txt', 'w', encoding='utf-8', errors='replace') as f: f.write( truncate_log( res[_LogKey_] ) )
+        # res[_LogKey_] = truncate_log( res[_LogKey_] )
+
+        # # Caution! Some of the strings in the result object may be unicode. Be robust to unicode in the log messages.
+        with codecs.open(setup.working_dir+'/.comparison.log.txt', 'w', encoding='utf-8', errors='replace') as f: f.write(res[_LogKey_])
+        truncate_results_logs(res)
+
+        print( 'Comparison finished with output:\n{}'.format( res[_LogKey_] ) )
+
+        with open(setup.working_dir+'/.comparison.results.json', 'w') as f: json.dump(res, f, sort_keys=True, indent=2)
+
+        #print( 'Comparison finished with results:\n{}'.format( json.dumps(res, sort_keys=True, indent=2) ) )
+        if 'summary' in res: print('Summary section:\n{}'.format( json.dumps(res['summary'], sort_keys=True, indent=2) ) )
+
+        print( f'Output results of this comparison saved to {working_dir_1}/.comparison.results.json\nComparison log saved into {working_dir_1}/.comparison.log.txt' )
+
+
+    else:
+        working_dir = setup.working_dir  #os.path.abspath( setup.config['results_root'] + f'/{platform["os"]}.{test}{options.suffix}' )
+
+        hpc_driver_name = setup.config['hpc_driver']
+        hpc_driver = None if hpc_driver_name in ['', 'none'] else eval(hpc_driver_name + '_HPC_Driver')(working_dir, setup.config, tracer=print, set_daemon_message=lambda x:None)
+
+        api_version = test_suite._api_version_ if hasattr(test_suite, '_api_version_') else ''
+
+        # if api_version < '1.0':
+        #     res = test_suite.run(test=test_name, rosetta_dir=os.path.abspath('../..'), working_dir=working_dir, platform=dict(Platform), jobs=Config.cpu_count, verbose=True, debug=Options.debug)
+        # else:
+
+        if api_version == '1.0': res = test_suite.run(test=test_name, repository_root=os.path.abspath('./..'), working_dir=working_dir, platform=dict(setup.platform), config=setup.config, hpc_driver=hpc_driver, verbose=True, debug=setup.debug)
+        else:
+            print(f'Test benchmark api_version={api_version} is not supported!'); sys.exit(1)
+
+        if not isinstance(res, dict): print(f'Test returned result of type {type(res)} while dict-like object was expected, please check that test-script have correct `return` statment! Terminating...'); sys.exit(1)
+
+        # Caution! Some of the strings in the result object may be unicode. Be robust to unicode in the log messages
+        with codecs.open(working_dir+'/.execution.log.txt', 'w', encoding='utf-8', errors='replace') as f: f.write( res[_LogKey_] )
+
+        # res[_LogKey_] = truncate_log( res[_LogKey_] )
+        truncate_results_logs(res)
+
+        if _DescriptionKey_ not in res: res[_DescriptionKey_] = test_description
+
+        if res[_StateKey_] not in _S_Values_: print( 'Warning!!! Test {} failed with unknow result code: {}'.format(test_name, res[_StateKey_]) )
+        else: print( f'Test {test} finished with output:\n{res[_LogKey_]}\n----------------------------------------------------------------\nState: {res[_StateKey_]!r} | ', end='')
+
+        # JSON by default serializes to an ascii-encoded format
+        with open(working_dir+'/.execution.results.json', 'w') as f: json.dump(res, f, sort_keys=True, indent=2)
+
+        print( f'Output and full log of this test saved to:\n{working_dir}/.execution.results.json\n{working_dir}/.execution.log.txt' )
+
+
+
+
+
+
+def main(args):
+    ''' Script to Run arbitrary Rosetta test
+    '''
+    parser = argparse.ArgumentParser(usage="Main testing script to run tests in the tests directory. "
+                                           "Use the --skip-compile to skip the build phase when testing locally. "
+                                           "Example Command: /benchmark.py -j2 integration.valgrind")
+
+    parser.add_argument('-j', '--jobs', default=0, type=int, help="Number of processors to use on when building. (default: use value from config file or 1)")
+
+    parser.add_argument('-m', '--memory', default=0, type=int, help="Amount of memory to use (default: use 2Gb per job")
+
+    parser.add_argument('--compiler', default=Platform['compiler'], help="Compiler to use")
+
+    #parser.add_argument('--python', default=('3.9' if Platform['os'] == 'mac' else '3.6'), help="Python interpreter to use")
+    parser.add_argument('--python', default=f'{sys.version_info.major}.{sys.version_info.minor}.s', help="Specify version of Python interpreter to use, for example '3.9'. If '.s' added to end of version string then use the same interpreter that was used to start this script. Default: '?.?.s'")
+
+    parser.add_argument("--extras", default='', help="Specify scons extras separated by ',': like --extras=mpi,static" )
+
+    parser.add_argument("--debug", action="store_true", dest="debug", default=False, help="Run specified test in debug mode (not with debug build!) this mean different things and depend on the test. Could be: skip the build phase, skip some of the test phases and so on. [off by default]" )
+
+    parser.add_argument("--suffix", default='', help="Specify ending suffix for test output dir. This is useful when you want to save test results in different dir for later comparison." )
+
+    parser.add_argument("--compare", nargs=2, help="Do not run the tests but instead compare previous results. Use --compare suffix1 suffix2" )
+
+    parser.add_argument("--config", default='benchmark.{os}.ini'.format(os=Platform['os']), action="store", help="Location of .ini file with additional options configuration. Optional.")
+
+    parser.add_argument("--skip-compile", dest='skip_compile', default=None, action="store_true", help="Skip the compilation phase. Assumes the binaries are already compiled locally.")
+
+    #parser.add_argument("--results-root", default=None, action="store", help="Location of `results` dir default is to use `./results`")
+
+    parser.add_argument("--setup", default=None, help="Specify JSON file with setup information. When this option supplied all other config and commandline options is ignored and auto-detection disable. Test, platform info will be gathered from provided JSON file. This option is designed to be used in daemon mode." )
+
+    parser.add_argument("--merge-head", default='HEAD', help="Specify SHA1/branch-name that will be used for `merge-head` value when simulating PR testing" )
+
+    parser.add_argument("--merge-base", default='origin/master', help="Specify SHA1/branch-name that will be used for `merge-base` value when simulating PR testing" )
+
+    parser.add_argument("--mount", action="append", help="Specify one of the mount points, like: --mount release_root:/some/path. This option could be used multiple times if needed" )
+
+
+    parser.add_argument('args', nargs=argparse.REMAINDER)
+
+    options = parser.parse_args(args=args[1:])
+
+    if any( [a.startswith('-') for a in options.args] ) :
+        print( '\nWARNING WARNING WARNING WARNING\n' )
+        print( '\tInterpreting', ' '.join(["'"+a+"'" for a in options.args if a.startswith('-')]), 'as test name(s), rather than as option(s).'  )
+        print( "\tTry moving it before any test name, if that's not what you want."  )
+        print( '\nWARNING WARNING WARNING WARNING\n'  )
+
+
+    if options.setup:
+        with open(options.setup) as f: setup = Setup( **json.load(f) )
+
+    else:
+        setup = setup_from_options(options)
+
+    run_test(setup)
+
+
+if __name__ == "__main__": main(sys.argv)

.rosetta-ci/benchmark.template.ini

@@ -0,0 +1,40 @@
+#
+# Benchmark script configuration file. Some of the tests require some system specific options to run. Please see benchmark.ini.template for list of available options.
+#
+
+[DEFAULT]
+
+[main] # additional config-options for various tests. All this fields will be pass as keys in 'config' function argument
+
+# how many jobs daemon can run on host machine (this is not related to HPC jobs)
+cpu_count = 24
+
+# how many memory in GB daemon can use on host machine (approximation, float)
+memory = 64
+
+# user name and email for user who submitted this test
+user_name  = Jane Roe
+user_email = jane.roe@university.edu
+
+# HPC Driver, might have one of the following values: MultiCore, Condor, Slurm or none if no HPC Driver should be configured
+hpc_driver = MultiCore
+
+# when running by daemons branch:revision will be set to appropriate values to represent currently checked version of main repository
+branch = unknown
+revision = 42
+
+# path to directory where test results will be stored
+results_root = ${_here_}/results
+
+release_root = ./results/_release_
+
+[slurm]
+# head-node host name, if specified will be used to submit jobs
+head_node =
+
+
+[mount]
+# list of key:path pairs that will be avalible as config.mounts during test run
+
+# path to releases, leave empty if release production should not be supported by this daemon
+release_root = ${_here_}/release

.rosetta-ci/hpc_drivers/__init__.py

@@ -0,0 +1,5 @@
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+from .multicore import MultiCore_HPC_Driver
+from .slurm     import Slurm_HPC_Driver

.rosetta-ci/hpc_drivers/base.py

@@ -0,0 +1,210 @@
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+import os, sys, subprocess, stat
+import time as time_module
+import signal as signal_module
+
+class NT:  # named tuple
+    def __init__(self, **entries): self.__dict__.update(entries)
+    def __repr__(self):
+        r = 'NT: |'
+        for i in dir(self):
+            if not i.startswith('__') and not isinstance(getattr(self, i), types.MethodType): r += '{} --> {}, '.format(i, getattr(self, i))
+        return r[:-2]+'|'
+
+
+
+class HPC_Exception(Exception):
+    def __init__(self, value): self.value = value
+    def __str__(self): return self.value
+
+
+
+def execute(message, command_line, return_='status', until_successes=False, terminate_on_failure=True, silent=False, silence_output=False, tracer=print):
+    if not silent: tracer(message);  tracer(command_line); sys.stdout.flush();
+    while True:
+
+        p = subprocess.Popen(command_line, bufsize=0, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        output, errors = p.communicate()
+
+        output = output + errors
+
+        output = output.decode(encoding="utf-8", errors="replace")
+
+        exit_code = p.returncode
+
+        if exit_code  and  not (silent or silence_output): tracer(output); sys.stdout.flush();
+
+        if exit_code and until_successes: pass  # Thats right - redability COUNT!
+        else: break
+
+        tracer( "Error while executing {}: {}\n".format(message, output) )
+        tracer("Sleeping 60s... then I will retry...")
+        sys.stdout.flush();
+        time.sleep(60)
+
+    if return_ == 'tuple': return(exit_code, output)
+
+    if exit_code and terminate_on_failure:
+        tracer("\nEncounter error while executing: " + command_line)
+        if return_==True: return True
+        else: print("\nEncounter error while executing: " + command_line + '\n' + output); sys.exit(1)
+
+    if return_ == 'output': return output
+    else: return False
+
+
+def Sleep(time_, message, dict_={}):
+    ''' Fancy sleep function '''
+    len_ = 0
+    for i in range(time_, 0, -1):
+        #print "Waiting for a new revision:%s... Sleeping...%d     \r" % (sc.revision, i),
+        msg = message.format( **dict(dict_, time_left=i) )
+        print( msg, end='' )
+        len_ = max(len_, len(msg))
+        sys.stdout.flush()
+        time_module.sleep(1)
+
+    print( ' '*len_ + '\r',  end='' ) # erazing sleep message
+
+
+# Abstract class for HPC job submission
+class HPC_Driver:
+    def __init__(self, working_dir, config, tracer=lambda x:None, set_daemon_message=lambda x:None):
+        self.working_dir = working_dir
+        self.config = config
+        self.cpu_usage  = 0.0  # cummulative cpu usage in hours
+        self.tracer     = tracer
+        self.set_daemon_message = set_daemon_message
+
+        self.cpu_count = self.config['cpu_count'] if type(config) == dict else self.config.getint('DEFAULT', 'cpu_count')
+
+        self.jobs = []  # list of all jobs currently running by this driver, Job class is driver depended, could be just int or something more complex
+
+        self.install_signal_handler()
+
+
+    def __del__(self):
+        self.remove_signal_handler()
+
+
+    def execute(self, executable, arguments, working_dir, log_dir=None, name='_no_name_', memory=256, time=24, shell_wrapper=False, block=True):
+        ''' Execute given command line on HPC cluster, must accumulate cpu hours in self.cpu_usage '''
+        if log_dir==None: log_dir=self.working_dir
+
+        if shell_wrapper:
+            shell_wrapper_sh = os.path.abspath(self.working_dir + '/hpc.{}.shell_wrapper.sh'.format(name))
+            with file(shell_wrapper_sh, 'w') as f: f.write('#!/bin/bash\n{} {}\n'.format(executable, arguments));  os.fchmod(f.fileno(), stat.S_IEXEC | stat.S_IREAD | stat.S_IWRITE)
+            executable, arguments = shell_wrapper_sh, ''
+
+        return self.submit_serial_hpc_job(name=name, executable=executable, arguments=arguments, working_dir=working_dir, log_dir=log_dir, jobs_to_queue=1, memory=memory, time=time, block=block, shell_wrapper=shell_wrapper)
+
+
+
+    @property
+    def number_of_cpu_per_node(self):
+        must_be_implemented_in_inherited_classes
+
+    @property
+    def maximum_number_of_mpi_cpu(self):
+        must_be_implemented_in_inherited_classes
+
+
+    def submit_hpc_job(self, name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory=512, time=12, block=True, shell_wrapper=False):
+        print('submit_hpc_job is DEPRECATED and will be removed in near future, please use submit_serial_hpc_job  instead!')
+        must_be_implemented_in_inherited_classes
+
+
+    def submit_serial_hpc_job(self, name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory=512, time=12, block=True, shell_wrapper=False):
+        must_be_implemented_in_inherited_classes
+
+
+    def submit_mpi_hpc_job(self, name, executable, arguments, working_dir, log_dir, memory=512, time=12, block=True, process_coefficient="1", requested_nodes=1, requested_processes_per_node=1):
+        ''' submit jobs as MPI job
+            process_coefficient should be string representing fraction of process to launch on each node, for example '3 / 4' will start only 75% of MPI process's on each node
+        '''
+        must_be_implemented_in_inherited_classes
+
+
+    def cancel_all_jobs(self):
+        ''' Cancel all HPC jobs known to this driver, use this as signal handler for script termination '''
+        for j in self.jobs: self.cancel_job(j)
+
+    def block_until(self, silent, fn, *args, **kwargs):
+        '''
+        **fn must have the driver as the first argument**
+        example:
+        def fn(driver):
+            jobs = list(driver.jobs)
+            jobs = [job for job in jobs if not driver.complete(job)]
+            if len(jobs) <= 8:
+                return False # stops sleeping
+            return True # continues sleeping
+
+        for x in range(100):
+            hpc_driver.submit_hpc_job(...)
+            hpc_driver.block_until(False, fn)
+        '''
+        while fn(self, *args, **kwargs):
+            sys.stdout.flush()
+            time_module.sleep(60)
+            if not silent:
+                Sleep(1, '"Waiting for HPC job(s) to finish, sleeping {time_left}s\r')
+
+    def wait_until_complete(self, jobs=None, callback=None, silent=False):
+        ''' Helper function, wait until given jobs list is finished, if no argument is given waits until all jobs known by driver is finished '''
+        jobs = jobs if jobs else self.jobs
+
+        while jobs:
+            for j in jobs[:]:
+                if self.complete(j): jobs.remove(j)
+
+            if jobs:
+                #total_cpu_queued  = sum( [j.jobs_queued  for j in jobs] )
+                #total_cpu_running = sum( [j.cpu_running for j in jobs] )
+                #self.set_daemon_message("Waiting for HPC job(s) to finish... [{} process(es) in queue, {} process(es) running]".format(total_cpu_queued, total_cpu_running) )
+                #self.tracer("Waiting for HPC job(s) [{} process(es) in queue, {} process(es) running]...  \r".format(total_cpu_queued, total_cpu_running), end='')
+                #print "Waiting for {} HPC jobs to finish... [{} jobs in queue, {} jobs running]... Sleeping 32s...     \r".format(total_cpu_queued, cpu_queued+cpu_running, cpu_running),
+
+                self.set_daemon_message("Waiting for HPC {} job(s) to finish...".format( len(jobs) ) )
+                #self.tracer("Waiting for HPC {} job(s) to finish...".format( len(jobs) ) )
+
+                sys.stdout.flush()
+
+                if callback: callback()
+
+                if silent: time_module.sleep(64*1)
+                else: Sleep(64, '"Waiting for HPC {n_jobs} job(s) to finish, sleeping {time_left}s    \r', dict(n_jobs=len(jobs)))
+
+
+
+    _signals_ = [signal_module.SIGINT, signal_module.SIGTERM, signal_module.SIGABRT]
+    def install_signal_handler(self):
+        def signal_handler(signal_, frame):
+            self.tracer('Recieved signal:{}... Canceling HPC jobs...'.format(signal_) )
+            self.cancel_all_jobs()
+            self.set_daemon_message( 'Remote daemon got terminated with signal:{}'.format(signal_) )
+            sys.exit(1)
+
+        for s in self._signals_: signal_module.signal(s, signal_handler)
+
+
+    def remove_signal_handler(self):  # do we really need this???
+        try:
+            for s in self._signals_: signal_module.signal(s, signal_module.SIG_DFL)
+            #print('remove_signal_handler: done!')
+
+        except TypeError:
+            #print('remove_signal_handler: interpreted terminating, skipping remove_signal_handler...')
+            pass
+
+
+    def cancel_job(self, job_id):
+        must_be_implemented_in_inherited_classes
+
+
+    def complete(self, job_id):
+        ''' Return job completion status. Return True if job complered and False otherwise
+        '''
+        must_be_implemented_in_inherited_classes

.rosetta-ci/hpc_drivers/multicore.py

@@ -0,0 +1,184 @@
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+import time as time_module
+import codecs
+import signal
+
+import os, sys
+
+try:
+    from .base import *
+
+except ImportError: # workaround for B2 back-end's
+    import imp
+    imp.load_source(__name__, '/'.join(__file__.split('/')[:-1]) + '/base.py')  # A bit of Python magic here, what we trying to say is this: from base import *, but path to base is calculated from our source location  # from base import HPC_Driver, execute, NT
+
+
+class MultiCore_HPC_Driver(HPC_Driver):
+
+    class JobID:
+        def __init__(self, pids=None):
+            self.pids = pids if pids else []
+
+
+        def __bool__(self): return bool(self.pids)
+
+
+        def __len__(self): return len(self.pids)
+
+
+        def add_pid(self, pid): self.pids.append(pid)
+
+
+        def remove_completed_pids(self):
+            for pid in self.pids[:]:
+                try:
+                    r = os.waitpid(pid, os.WNOHANG)
+                    if r == (pid, 0): self.pids.remove(pid)  # process have ended without error
+                    elif r[0] == pid :  # process ended but with error, special case we will have to wait for all process to terminate and call system exit.
+                        #self.cancel_job()
+                        #sys.exit(1)
+                        self.pids.remove(pid)
+                        print('ERROR: Some of the HPC jobs terminated abnormally! Please see HPC logs for details.')
+
+                except ChildProcessError: self.pids.remove(pid)
+
+
+        def cancel(self):
+            for pid in self.pids:
+                try:
+                    os.killpg(os.getpgid(pid), signal.SIGKILL)
+                except ChildProcessError: pass
+
+            self.pids = []
+
+
+
+    def __init__(self, *args, **kwds):
+        HPC_Driver.__init__(self, *args, **kwds)
+        #print(f'MultiCore_HPC_Driver: cpu_count: {self.cpu_count}')
+
+
+    def remove_completed_jobs(self):
+        for job in self.jobs[:]:  # Need to make a copy so we don't modify a list we're iterating over
+            job.remove_completed_pids()
+            if not job: self.jobs.remove(job)
+
+
+    @property
+    def process_count(self):
+        ''' return number of processes that currently ran by this driver instance
+        '''
+        return sum( map(len, self.jobs) )
+
+
+    def submit_hpc_job(self, name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory=512, time=12, block=True, shell_wrapper=False):
+        print('submit_hpc_job is DEPRECATED and will be removed in near future, please use submit_serial_hpc_job  instead!')
+        return self.submit_serial_hpc_job(name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory, time, block, shell_wrapper)
+
+
+    def submit_serial_hpc_job(self, name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory=512, time=12, block=True, shell_wrapper=False):
+        cpu_usage = -time_module.time()/60./60.
+
+        if shell_wrapper:
+            shell_wrapper_sh = os.path.abspath(self.working_dir + f'/hpc.{name}.shell_wrapper.sh')
+            with open(shell_wrapper_sh, 'w') as f: f.write('#!/bin/bash\n{} {}\n'.format(executable, arguments));  os.fchmod(f.fileno(), stat.S_IEXEC | stat.S_IREAD | stat.S_IWRITE)
+            executable, arguments = shell_wrapper_sh, ''
+
+        def mfork():
+            ''' Check if number of child process is below cpu_count. And if it is - fork the new pocees and return its pid.
+            '''
+            while self.process_count >= self.cpu_count:
+                self.remove_completed_jobs()
+                if self.process_count >= self.cpu_count: time_module.sleep(.5)
+
+            sys.stdout.flush()
+            pid = os.fork()
+            # appending at caller level insted  if pid: self.jobs.append(pid) # We are parent!
+            return pid
+
+        current_job = self.JobID()
+        process = 0
+        for i in range(jobs_to_queue):
+
+            pid = mfork()
+            if not pid: # we are child process
+                command_line = 'cd {} && {} {}'.format(working_dir, executable, arguments.format(process=process) )
+                exit_code, log = execute('Running job {}.{}...'.format(name, i), command_line, tracer=self.tracer, return_='tuple')
+                with codecs.open(log_dir+'/.hpc.{name}.{i:02d}.log'.format(**vars()), 'w', encoding='utf-8', errors='replace') as f:
+                    f.write(command_line+'\n'+log)
+                    if exit_code:
+                        error_report = f'\n\n{command_line}\nERROR: PROCESS {name}.{i:02d} TERMINATED WITH NON-ZERO-EXIT-CODE {exit_code}!\n'
+                        f.write(error_report)
+                        print(log, error_report)
+
+                sys.exit(0)
+
+            else: # we are parent!
+                current_job.add_pid(pid)
+                # Need to potentially re-add to list, as remove_completed_jobs() might trim it.
+                if current_job not in self.jobs: self.jobs.append(current_job)
+
+            process += 1
+
+        if block:
+            #for p in all_queued_jobs: os.waitpid(p, 0)  # waiting for all child process to termintate...
+
+            self.wait_until_complete(current_job)
+            self.remove_completed_jobs()
+
+            cpu_usage += time_module.time()/60./60.
+            self.cpu_usage += cpu_usage * jobs_to_queue  # approximation...
+
+            current_job = self.JobID()
+
+        return current_job
+
+
+    @property
+    def number_of_cpu_per_node(self): return self.cpu_count
+
+
+    @property
+    def maximum_number_of_mpi_cpu(self): return self.cpu_count
+
+
+    def submit_mpi_hpc_job(self, name, executable, arguments, working_dir, log_dir, memory=512, time=12, block=True, process_coefficient="1", requested_nodes=1, requested_processes_per_node=1):
+
+        if requested_nodes > 1:
+            print( "WARNING: " + str( requested_nodes ) + " nodes were requested, but we're running locally, so only 1 node will be used." )
+
+        if requested_processes_per_node > self.cpu_count:
+            print( "WARNING: " + str(requested_processes_per_node) + " processes were requested, but I only have " + str(self.cpu_count) + " CPUs.  Will launch " + str(self.cpu_count) + " processes."  )
+        actual_processes = min( requested_processes_per_node, self.cpu_count )
+
+        cpu_usage = -time_module.time()/60./60.
+
+        arguments = arguments.format(process=0)
+
+        command_line = f'cd {working_dir} && mpirun -np {actual_processes} {executable} {arguments}'
+        log = execute(f'Running job {name}...', command_line, tracer=self.tracer, return_='output')
+        with codecs.open(log_dir+'/.hpc.{name}.log'.format(**vars()), 'w', encoding='utf-8', errors='replace') as f: f.write(command_line+'\n'+log)
+
+        cpu_usage += time_module.time()/60./60.
+        self.cpu_usage += cpu_usage * actual_processes  # approximation...
+
+        # return None - we do not return anything from this version of submit which imply returning None which in turn will be treated as job-id for already finished job
+
+
+    def complete(self, job_id):
+        ''' Return job completion status. Return True if job completed and False otherwise
+        '''
+        self.remove_completed_jobs()
+        return job_id not in self.jobs
+
+
+    def cancel_job(self, job):
+        job.cancel();
+        if job in self.jobs:
+            self.jobs.remove(job)
+
+
+    def __repr__(self):
+        return 'MultiCore_HPC_Driver<>'

.rosetta-ci/hpc_drivers/slurm.py

@@ -0,0 +1,176 @@
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+import os, sys, time, collections, math
+import stat as stat_module
+
+
+try:
+    from .base import *
+
+except ImportError: # workaround for B2 back-end's
+    import imp
+    imp.load_source(__name__, '/'.join(__file__.split('/')[:-1]) + '/base.py')  # A bit of Python magic here, what we trying to say is this: from base import *, but path to base is calculated from our source location  # from base import HPC_Driver, execute, NT
+
+
+_T_slurm_array_job_template_ = '''\
+#!/bin/bash
+#
+#SBATCH --job-name={name}
+#SBATCH --output={log_dir}/.hpc.%x.%a.output
+#
+#SBATCH --time={time}:00
+#SBATCH --mem-per-cpu={memory}M
+#SBATCH --chdir={working_dir}
+#
+#SBATCH --array=1-{jobs_to_queue}
+
+srun {executable} {arguments}
+'''
+
+_T_slurm_mpi_job_template_ = '''\
+#!/bin/bash
+#
+#SBATCH --job-name={name}
+#SBATCH --output={log_dir}/.hpc.%x.output
+#
+#SBATCH --time={time}:00
+#SBATCH --mem-per-cpu={memory}M
+#SBATCH --chdir={working_dir}
+#
+#SBATCH --ntasks={ntasks}
+
+mpirun {executable} {arguments}
+'''
+
+class Slurm_HPC_Driver(HPC_Driver):
+    def head_node_execute(self, message, command_line, *args, **kwargs):
+        head_node = self.config['slurm'].get('head_node')
+
+        command_line, host = (f"ssh {head_node} cd `pwd` '&& {command_line}'", head_node) if head_node else (command_line, 'localhost')
+        return execute(f'Executiong on {host}: {message}' if message else '', command_line, *args, **kwargs)
+
+
+    # NodeGroup = collections.namedtuple('NodeGroup', 'nodes cores')
+
+    # @property
+    # def mpi_topology(self):
+    #     ''' return list of NodeGroup's
+    #     '''
+    #     pass
+
+
+    # @property
+    # def number_of_cpu_per_node(self): return int( self.config['condor']['mpi_cpu_per_node'] )
+
+    # @property
+    # def maximum_number_of_mpi_cpu(self):
+    #     return self.number_of_cpu_per_node * int( self.config['condor']['mpi_maximum_number_of_nodes'] )
+
+
+    # def complete(self, condor_job_id):
+    #     ''' Return job completion status. Note that single hpc_job may contatin inner list of individual HPC jobs, True should be return if they all run in to completion.
+    #     '''
+
+    #     execute('Releasing condor jobs...', 'condor_release $USER', return_='tuple')
+
+    #     s = execute('', 'condor_q $USER | grep $USER | grep {}'.format(condor_job_id), return_='output', terminate_on_failure=False).replace(' ', '').replace('\n', '')
+    #     if s: return False
+
+    #         # #setDaemonStatusAndPing('[Job #%s] Running... %s condor job(s) in queue...' % (self.id, len(s.split('\n') ) ) )
+    #         # n_jobs = len(s.split('\n'))
+    #         # s, o = execute('', 'condor_userprio -all | grep $USER@', return_='tuple')
+    #         # if s == 0:
+    #         #     jobs_running = o.split()
+    #         #     jobs_running = 'XX' if len(jobs_running) < 4 else jobs_running[4]
+    #         #     self.set_daemon_message("Waiting for condor to finish HPC jobs... [{} jobs in HPC-Queue, {} CPU's used]".format(n_jobs, jobs_running) )
+    #         #     print "{} condor jobs in queue... Sleeping 32s...    \r".format(n_jobs),
+    #         # sys.stdout.flush()
+    #         # time.sleep(32)
+    #     else:
+
+    #         #self.tracer('Waiting for condor to finish the jobs... DONE')
+    #         self.jobs.remove(condor_job_id)
+    #         self.cpu_usage += self.get_condor_accumulated_usage()
+    #         return True  # jobs already finished, we return empty list to prevent double counting of cpu_usage
+
+
+    def complete(self, slurm_job_id):
+        ''' Return True if job with given id is complete
+        '''
+
+        s = self.head_node_execute('', f'squeue -j {slurm_job_id} --noheader', return_='output', terminate_on_failure=False, silent=True)
+        if s: return False
+        else:
+            #self.tracer('Waiting for condor to finish the jobs... DONE')
+            self.jobs.remove(slurm_job_id)
+            return True  # jobs already finished, we return empty list to prevent double counting of cpu_usage
+
+
+    def cancel_job(self, slurm_job_id):
+        self.head_node_execute(f'Slurm_HPC_Driver.canceling job {slurm_job_id}...', f'scancel {slurm_job_id}', terminate_on_failure=False)
+
+
+    # def submit_hpc_job(self, name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory=512, time=12, block=True, shell_wrapper=False):
+    #     print('submit_hpc_job is DEPRECATED and will be removed in near future, please use submit_serial_hpc_job  instead!')
+    #     return self.submit_serial_hpc_job(name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory, time, block, shell_wrapper)
+
+
+    def submit_serial_hpc_job(self, name, executable, arguments, working_dir, jobs_to_queue, log_dir, memory=512, time=12, block=True, shell_wrapper=False):
+
+        arguments = arguments.format(process='%a') # %a is SLURM array index
+        time = int( math.ceil(time*60) )
+
+        if shell_wrapper:
+            shell_wrapper_sh = os.path.abspath(self.working_dir + f'/hpc.{name}.shell_wrapper.sh')
+            with open(shell_wrapper_sh, 'w') as f: f.write('#!/bin/bash\n{} {}\n'.format(executable, arguments));  os.fchmod(f.fileno(), stat.S_IEXEC | stat.S_IREAD | stat.S_IWRITE)
+            executable, arguments = shell_wrapper_sh, ''
+
+        slurm_file = working_dir + f'/.hpc.{name}.slurm'
+
+        with open(slurm_file, 'w') as f: f.write( _T_slurm_array_job_template_.format( **vars() ) )
+
+
+        slurm_job_id = self.head_node_execute('Submitting SLURM array job...', f'cd {self.working_dir} && sbatch {slurm_file}',
+                                              tracer=self.tracer, return_='output'
+                                              ).split()[-1] # expecting something like `Submitted batch job 6122` in output
+
+
+        self.jobs.append(slurm_job_id)
+
+        if block:
+            self.wait_until_complete( [slurm_job_id] )
+            return None
+
+        else: return slurm_job_id
+
+
+
+
+
+    def submit_mpi_hpc_job(self, name, executable, arguments, working_dir, log_dir, ntasks, memory=512, time=12, block=True, shell_wrapper=False):
+        ''' submit jobs as MPI job
+        '''
+        arguments = arguments.format(process='0')
+        time = int( math.ceil(time*60) )
+
+        if shell_wrapper:
+            shell_wrapper_sh = os.path.abspath(self.working_dir + f'/hpc.{name}.shell_wrapper.sh')
+            with open(shell_wrapper_sh, 'w') as f: f.write('#!/bin/bash\n{} {}\n'.format(executable, arguments));  os.fchmod(f.fileno(), stat.S_IEXEC | stat.S_IREAD | stat.S_IWRITE)
+            executable, arguments = shell_wrapper_sh, ''
+
+        slurm_file = working_dir + f'/.hpc.{name}.slurm'
+
+        with open(slurm_file, 'w') as f: f.write( _T_slurm_mpi_job_template_.format( **vars() ) )
+
+        slurm_job_id = self.head_node_execute('Submitting SLURM mpi job...', f'cd {self.working_dir} && sbatch {slurm_file}',
+                                              tracer=self.tracer, return_='output'
+                                              ).split()[-1] # expecting something like `Submitted batch job 6122` in output
+
+        self.jobs.append(slurm_job_id)
+
+        if block:
+            self.wait_until_complete( [slurm_job_id] )
+            return None
+
+        else: return slurm_job_id

.rosetta-ci/test-sets.yaml

@@ -0,0 +1,65 @@
+# map platform-string → platform definiton
+platforms:
+  ubuntu-20.04.gcc:
+    os: ubuntu-20.04
+    compiler: gcc
+    python: '3.9'
+
+  ubuntu-20.04.clang:
+    os: ubuntu-20.04
+    compiler: clang
+    python: '3.9'
+
+
+# map of test-set-name → tests
+test-sets:
+  main:
+    - ubuntu-20.04.clang.rfd
+
+  python:
+    - ubuntu-20.04.gcc.self.python
+    - ubuntu-20.04.clang.self.python
+
+  self:
+    - ubuntu-20.04.gcc.self.state
+    - ubuntu-20.04.gcc.self.subtests
+    - ubuntu-20.04.gcc.self.release
+
+
+
+# map of GitHub-label → [test-set]
+github-label-test-sets:
+  00 main: [main]
+  10 self: [self]
+  16 python: [python]
+
+
+# map of submit-page-category → tests
+# tests that does not get assigned will be automatically displayed in 'other' category
+category-tests:
+  main:
+    - rfd
+
+  self:
+    - self.state
+    - self.subtests
+    - self.release
+    - self.python
+
+
+# map branch → test-set to
+# specify list of tests that should be applied by-default during testing of each new commits to specific branch
+branch-test-sets:
+  main: [main]
+  benchmark: [main, python]
+
+
+# map branch → test-sets for pull-request's
+# specify which test-sets should be scheduled for PR's by-default (ie in addition to GH labels applied)
+# use empty branch name to specify defult value for (ie any branch not explicitly listed)
+pull-request-branch-test-sets:
+  # specific test sets for benchmark branch
+  benchmark: ['main', 'python']
+
+  # default, will apply to PR's to any other branch
+  '': ['main']

.rosetta-ci/tests/__init__.py

@@ -0,0 +1,765 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+# (c) Copyright Rosetta Commons Member Institutions.
+# (c) This file is part of the Rosetta software suite and is made available under license.
+# (c) The Rosetta software is developed by the contributing members of the Rosetta Commons.
+# (c) For more information, see http://www.rosettacommons.org. Questions about this can be
+# (c) addressed to University of Washington CoMotion, email: license@uw.edu.
+
+## @file   tests/__init__.py
+## @brief  Common constats and types for all test types
+## @author Sergey Lyskov
+
+import os, time, sys, shutil, codecs, urllib.request, imp, subprocess, json, hashlib  # urllib.error, urllib.parse,
+import platform as  platform_module
+import types as types_module
+
+# ⚔ do not change wording below, it have to stay in sync with upstream (up to benchmark-model).
+# Copied from benchmark-model, standard state code's for tests results.
+
+__all__ = ['execute',
+           '_S_Values_', '_S_draft_', '_S_queued_', '_S_running_', '_S_passed_', '_S_failed_', '_S_build_failed_', '_S_script_failed_',
+           '_StateKey_', '_ResultsKey_', '_LogKey_', '_DescriptionKey_', '_TestsKey_',
+           '_multi_step_config_', '_multi_step_error_', '_multi_step_result_',
+           'to_bytes',
+]
+
+_S_draft_                 = 'draft'
+_S_queued_                = 'queued'
+_S_running_               = 'running'
+_S_passed_                = 'passed'
+_S_failed_                = 'failed'
+_S_build_failed_          = 'build failed'
+_S_script_failed_         = 'script failed'
+_S_queued_for_comparison_ = 'queued for comparison'
+
+_S_Values_ = [_S_draft_, _S_queued_, _S_running_, _S_passed_, _S_failed_, _S_build_failed_, _S_script_failed_, _S_queued_for_comparison_]
+
+_IgnoreKey_      = 'ignore'
+_StateKey_       = 'state'
+_ResultsKey_     = 'results'
+_LogKey_         = 'log'
+_DescriptionKey_ = 'description'
+_TestsKey_       = 'tests'
+_SummaryKey_     = 'summary'
+_FailedKey_      = 'failed'
+_TotalKey_       = 'total'
+_PlotsKey_       = 'plots'
+_FailedTestsKey_ = 'failed_tests'
+_HtmlKey_        = 'html'
+
+# file names for multi-step test files
+_multi_step_config_ = 'config.json'
+_multi_step_error_  = 'error.json'
+_multi_step_result_ = 'result.json'
+
+PyRosetta_unix_memory_requirement_per_cpu = 6  # Memory per sub-process in Gb's
+PyRosetta_unix_unit_test_memory_requirement_per_cpu = 3.0  # Memory per sub-process in Gb's for running PyRosetta unit tests
+
+# Commands to run all the scripts needed for setting up Rosetta compiles. (Run from main/source directory)
+PRE_COMPILE_SETUP_SCRIPTS = [ "./update_options.sh", "./update_submodules.sh", "./update_ResidueType_enum_files.sh", "python version.py" ]
+
+DEFAULT_PYTHON_VERSION='3.9'
+
+# Standard funtions and classes below ---------------------------------------------------------------------------------
+
+class BenchmarkError(Exception):
+    def __init__(self, value): self.value = value
+    def __repr__(self): return self.value
+    def __str__(self): return self.value
+
+
+class NT:  # named tuple
+    def __init__(self, **entries): self.__dict__.update(entries)
+    def __repr__(self):
+        r = 'NT: |'
+        for i in dir(self):
+            print(i)
+            if not i.startswith('__') and i != '_as_dict' and not isinstance(getattr(self, i), types_module.MethodType): r += '%s --> %s, ' % (i, getattr(self, i))
+        return r[:-2]+'|'
+
+    @property
+    def _as_dict(self):
+        return { a: getattr(self, a) for a in dir(self) if not a.startswith('__') and a != '_as_dict' and not isinstance(getattr(self, a), types_module.MethodType)}
+
+
+def Tracer(verbose=False):
+    return print if verbose else lambda x: None
+
+
+def to_unicode(b):
+    ''' Conver bytes to string and handle the errors. If argument is already in string - do nothing
+    '''
+    #return b if type(b) == unicode else unicode(b, 'utf-8', errors='replace')
+    return b if type(b) == str else str(b, 'utf-8', errors='backslashreplace')
+
+
+def to_bytes(u):
+    ''' Conver string to bytes and handle the errors. If argument is already of type bytes - do nothing
+    '''
+    return u if type(u) == bytes else u.encode('utf-8', errors='backslashreplace')
+
+
+''' Python-2 version
+def execute(message, commandline, return_=False, until_successes=False, terminate_on_failure=True, add_message_and_command_line_to_output=False):
+    message, commandline = to_unicode(message), to_unicode(commandline)
+
+    TR = Tracer()
+    TR(message);  TR(commandline)
+    while True:
+        (res, output) = commands.getstatusoutput(commandline)
+        # Subprocess results will always be a bytes-string.
+        # Probably ASCII, but may have some Unicode characters.
+        # A UTF-8 decode will probably get decent results 99% of the time
+        # and the replace option will gracefully handle the rest.
+        output = to_unicode(output)
+
+        TR(output)
+
+        if res and until_successes: pass  # Thats right - redability COUNT!
+        else: break
+
+        print( "Error while executing %s: %s\n" % (message, output) )
+        print( "Sleeping 60s... then I will retry..." )
+        time.sleep(60)
+
+    if add_message_and_command_line_to_output: output = message + '\nCommand line: ' + commandline + '\n' + output
+
+    if return_ == 'tuple': return(res, output)
+
+    if res and terminate_on_failure:
+        TR("\nEncounter error while executing: " + commandline)
+        if return_==True: return res
+        else:
+            print("\nEncounter error while executing: " + commandline + '\n' + output)
+            raise BenchmarkError("\nEncounter error while executing: " + commandline + '\n' + output)
+
+    if return_ == 'output': return output
+    else: return res
+'''
+
+def execute_through_subprocess(command_line):
+    # exit_code, output = subprocess.getstatusoutput(command_line)
+
+    # p = subprocess.Popen(command_line, bufsize=0, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    # output, errors = p.communicate()
+    # output = (output + errors).decode(encoding='utf-8', errors='backslashreplace')
+    # exit_code = p.returncode
+
+    # previous 'main' version based on subprocess module. Main issue that output of segfaults will not be captured since they generated by shell
+    p = subprocess.Popen(command_line, bufsize=0, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    output, errors = p.communicate()
+    # output = output + errors # ← we redirected stderr into same pipe as stdcout so errors is None, - no need to concatenate
+    output = output.decode(encoding='utf-8', errors='backslashreplace')
+    exit_code = p.returncode
+
+    return exit_code, output
+
+
+def execute_through_pexpect(command_line):
+    import pexpect
+
+    child = pexpect.spawn('/bin/bash', ['-c', command_line])
+    child.expect(pexpect.EOF)
+    output = child.before.decode(encoding='utf-8', errors='backslashreplace')
+    child.close()
+    exit_code = child.signalstatus or child.exitstatus
+
+    return exit_code, output
+
+
+def execute_through_pty(command_line):
+    import pty, select
+
+    if sys.platform == "darwin":
+
+        master, slave = pty.openpty()
+        p = subprocess.Popen(command_line, shell=True, stdout=slave, stdin=slave,
+                             stderr=subprocess.STDOUT, close_fds=True)
+
+        buffer = []
+        while True:
+            try:
+                if select.select([master], [], [], 0.2)[0]:  # has something to read
+                    data = os.read(master, 1 << 22)
+                    if data: buffer.append(data)
+
+                elif (p.poll() is not None)  and  (not select.select([master], [], [], 0.2)[0] ): break  # process is finished and output buffer if fully read
+
+            except OSError: break  # OSError will be raised when child process close PTY descriptior
+
+        output = b''.join(buffer).decode(encoding='utf-8', errors='backslashreplace')
+
+        os.close(master)
+        os.close(slave)
+
+        p.wait()
+        exit_code = p.returncode
+
+        '''
+        buffer = []
+        while True:
+            if select.select([master], [], [], 0.2)[0]:  # has something to read
+                data = os.read(master, 1 << 22)
+                if data: buffer.append(data)
+                # else: break  # # EOF - well, technically process _should_ be finished here...
+
+            # elif time.sleep(1) or (p.poll() is not None): # process is finished (sleep here is intentional to trigger race condition, see solution for this on the next few lines)
+            #     assert not select.select([master], [], [], 0.2)[0]  # should be nothing left to read...
+            #     break
+
+            elif (p.poll() is not None)  and  (not select.select([master], [], [], 0.2)[0] ): break  # process is finished and output buffer if fully read
+
+        assert not select.select([master], [], [], 0.2)[0]  # should be nothing left to read...
+
+        os.close(slave)
+        os.close(master)
+
+        output = b''.join(buffer).decode(encoding='utf-8', errors='backslashreplace')
+        exit_code = p.returncode
+        '''
+
+    else:
+
+        master, slave = pty.openpty()
+        p = subprocess.Popen(command_line, shell=True, stdout=slave, stdin=slave,
+                             stderr=subprocess.STDOUT, close_fds=True)
+
+        os.close(slave)
+
+        buffer = []
+        while True:
+            try:
+                data = os.read(master, 1 << 22)
+                if data: buffer.append(data)
+            except OSError: break  # OSError will be raised when child process close PTY descriptior
+
+        output = b''.join(buffer).decode(encoding='utf-8', errors='backslashreplace')
+
+        os.close(master)
+
+        p.wait()
+        exit_code = p.returncode
+
+    return exit_code, output
+
+
+
+def execute(message, command_line, return_='status', until_successes=False, terminate_on_failure=True, silent=False, silence_output=False, silence_output_on_errors=False, add_message_and_command_line_to_output=False):
+    if not silent: print(message);  print(command_line); sys.stdout.flush();
+    while True:
+
+        #exit_code, output = execute_through_subprocess(command_line)
+        #exit_code, output = execute_through_pexpect(command_line)
+        exit_code, output = execute_through_pty(command_line)
+
+        if (exit_code  and  not silence_output_on_errors) or  not (silent or silence_output): print(output); sys.stdout.flush();
+
+        if exit_code and until_successes: pass  # Thats right - redability COUNT!
+        else: break
+
+        print( "Error while executing {}: {}\n".format(message, output) )
+        print("Sleeping 60s... then I will retry...")
+        sys.stdout.flush();
+        time.sleep(60)
+
+    if add_message_and_command_line_to_output: output = message + '\nCommand line: ' + command_line + '\n' + output
+
+    if return_ == 'tuple'  or  return_ == tuple: return(exit_code, output)
+
+    if exit_code and terminate_on_failure:
+        print("\nEncounter error while executing: " + command_line)
+        if return_==True: return True
+        else:
+            print('\nEncounter error while executing: ' + command_line + '\n' + output);
+            raise BenchmarkError('\nEncounter error while executing: ' + command_line + '\n' + output)
+
+    if return_ == 'output': return output
+    else: return exit_code
+
+
+def parallel_execute(name, jobs, rosetta_dir, working_dir, cpu_count, time=16):
+    ''' Execute command line in parallel on local host
+        time specifies the upper limit for cpu-usage runtime (in minutes) for any one process in the parallel execution.
+
+        jobs should be dict with following structure:
+        {
+            'job-string-id-1’: command_line-1,
+            'job-string-id-2’: command_line-2,
+            ...
+        }
+
+        return: dict with jobs-id's as keys and value as dict with 'output' and 'result' keys:
+        {
+            "job-string-id-1": {
+                "output": "stdout + stdderr output of command_line-1",
+                "result": <integer exit code for command_line-1>
+            },
+            "c2": {
+                "output": "stdout + stdderr output of command_line-2",
+                "result": <integer exit code for command_line-2>
+            },
+            ...
+        }
+    '''
+    job_file_name = working_dir + '/' + name
+    with open(job_file_name + '.json', 'w') as f: json.dump(jobs, f, sort_keys=True, indent=2) # JSON handles unicode internally
+    if time is not None:
+        allowed_time = int(time*60)
+        ulimit_command = f'ulimit -t {allowed_time} && '
+    else:
+        ulimit_command = ''
+    command = f'cd {working_dir} && ' + ulimit_command + f'{rosetta_dir}/tests/benchmark/util/parallel.py -j{cpu_count} {job_file_name}.json'
+    execute("Running {} in parallel with {} CPU's...".format(name, cpu_count), command )
+
+    with open(job_file_name+'.results.json') as f: return json.load(f)
+
+
+def calculate_unique_prefix_path(platform, config):
+    ''' calculate path for prefix location that is unique for this machine and OS
+    '''
+    hostname = os.uname()[1]
+    return config['prefix'] + '/' + hostname + '/' + platform['os']
+
+
+def get_python_include_and_lib(python):
+    ''' calculate python include dir and lib dir from given python executable path
+    '''
+    #python = os.path.realpath(python)
+    python_bin_dir = python.rpartition('/')[0]
+    python_config = f'{python} {python}-config' if python.endswith('2.7') else f'{python}-config'
+
+    #if not os.path.isfile(python_config): python_config = python_bin_dir + '/python-config'
+
+    info = execute('Getting python configuration info...', f'unset __PYVENV_LAUNCHER__ && cd {python_bin_dir} && PATH=.:$PATH && {python_config} --prefix --includes', return_='output').replace('\r', '').split('\n')  # Python-3 only: --abiflags
+    python_prefix = info[0]
+    python_include_dir = info[1].split()[0][len('-I'):]
+    python_lib_dir = python_prefix + '/lib'
+    #python_abi_suffix = info[2]
+    #print(python_include_dir, python_lib_dir)
+
+    return NT(python_include_dir=python_include_dir, python_lib_dir=python_lib_dir)
+
+
+def local_open_ssl_install(prefix, build_prefix, jobs):
+    ''' install OpenSSL at given prefix, return url of source archive
+    '''
+    #with tempfile.TemporaryDirectory('open_ssl_build', dir=prefix) as build_prefix:
+
+    url = 'https://www.openssl.org/source/openssl-1.1.1b.tar.gz'
+    #url = 'https://www.openssl.org/source/openssl-3.0.0.tar.gz'
+
+
+    archive = build_prefix + '/' + url.split('/')[-1]
+    build_dir = archive.rpartition('.tar.gz')[0]
+    if os.path.isdir(build_dir): shutil.rmtree(build_dir)
+
+    with open(archive, 'wb') as f:
+        response = urllib.request.urlopen(url)
+        f.write( response.read() )
+
+    execute('Unpacking {}'.format(archive), 'cd {build_prefix} && tar -xvzf {archive}'.format(**vars()) )
+
+    execute('Configuring...', f'cd {build_dir} && ./config --prefix={prefix}')
+    execute('Building...',    f'cd {build_dir} && make -j{jobs}')
+    execute('Installing...',  f'cd {build_dir} && make -j{jobs} install')
+
+    return url
+
+
+def remove_pip_and_easy_install(prefix_root_path):
+    ''' remove `pip` and `easy_install` executable from given Python / virtual-environments install
+    '''
+    for f in os.listdir(prefix_root_path + '/bin'):  # removing all pip's and easy_install's to make sure that environment is immutable
+        for p in ['pip', 'easy_install']:
+            if f.startswith(p): os.remove(prefix_root_path + '/bin/' + f)
+
+
+
+def local_python_install(platform, config):
+    ''' Perform local install of given Python version and return path-to-python-interpreter, python_include_dir, python_lib_dir
+        If previous install is detected skip installiation.
+        Provided Python install will _persistent_ and _immutable_
+    '''
+    jobs = config['cpu_count']
+    compiler, cpp_compiler = ('clang', 'clang++') if platform['os'] == 'mac' else ('gcc', 'g++')  # disregarding platform compiler setting and instead use default compiler for platform
+
+    python_version = platform.get('python', DEFAULT_PYTHON_VERSION)
+
+    if python_version.endswith('.s'):
+        assert python_version == f'{sys.version_info.major}.{sys.version_info.minor}.s'
+        #root = executable.rpartition('/bin/python')[0]
+        h = hashlib.md5(); h.update( (sys.executable + sys.version).encode('utf-8', errors='backslashreplace') ); hash = h.hexdigest()
+        return NT(
+            python = sys.executable,
+            root = None,
+            python_include_dir = None,
+            python_lib_dir = None,
+            version = python_version,
+            url = None,
+            platform = platform,
+            config = config,
+            hash = hash,
+        )
+
+    # deprecated, no longer needed
+    # python_version = {'python2'   : '2.7',
+    #                   'python2.7' : '2.7',
+    #                   'python3'   : '3.5',
+    # }.get(python_version, python_version)
+
+    # for security reasons we only allow installs for version listed here with hand-coded URL's
+    python_sources = {
+        '2.7' : 'https://www.python.org/ftp/python/2.7.18/Python-2.7.18.tgz',
+
+        '3.5'  : 'https://www.python.org/ftp/python/3.5.9/Python-3.5.9.tgz',
+        '3.6'  : 'https://www.python.org/ftp/python/3.6.15/Python-3.6.15.tgz',
+        '3.7'  : 'https://www.python.org/ftp/python/3.7.14/Python-3.7.14.tgz',
+        '3.8'  : 'https://www.python.org/ftp/python/3.8.14/Python-3.8.14.tgz',
+        '3.9'  : 'https://www.python.org/ftp/python/3.9.14/Python-3.9.14.tgz',
+        '3.10' : 'https://www.python.org/ftp/python/3.10.10/Python-3.10.10.tgz',
+        '3.11' : 'https://www.python.org/ftp/python/3.11.2/Python-3.11.2.tgz',
+    }
+
+    # map of env -> ('shell-code-before ./configure', 'extra-arguments-for-configure')
+    extras = {
+        #('mac',) :          ('__PYVENV_LAUNCHER__="" MACOSX_DEPLOYMENT_TARGET={}'.format(platform_module.mac_ver()[0]), ''),  # __PYVENV_LAUNCHER__ now used by-default for all platform installs
+        ('mac',) :          ('MACOSX_DEPLOYMENT_TARGET={}'.format(platform_module.mac_ver()[0]), ''),
+        ('linux',  '2.7') : ('', '--enable-unicode=ucs4'),
+        ('ubuntu', '2.7') : ('', '--enable-unicode=ucs4'),
+    }
+
+    #packages = '' if (python_version[0] == '2' or  python_version == '3.5' ) and  platform['os'] == 'mac' else 'pip setuptools wheel' # 2.7 is now deprecated on Mac so some packages could not be installed
+    packages = 'setuptools'
+
+    url = python_sources[python_version]
+
+    extra = extras.get( (platform['os'],)  , ('', '') )
+    extra = extras.get( (platform['os'], python_version) , extra)
+
+    extra = ('unset __PYVENV_LAUNCHER__ && ' + extra[0], extra[1])
+
+    options = '--with-ensurepip' #'--without-ensurepip'
+    signature = f'v1.5.1 url: {url}\noptions: {options}\ncompiler: {compiler}\nextra: {extra}\npackages: {packages}\n'
+
+    h = hashlib.md5(); h.update( signature.encode('utf-8', errors='backslashreplace') ); hash = h.hexdigest()
+
+    root = calculate_unique_prefix_path(platform, config) + '/python-' + python_version + '.' +  compiler + '/' + hash
+
+    signature_file_name = root + '/.signature'
+
+    #activate   = root + '/bin/activate'
+    executable = root + '/bin/python' + python_version
+
+    # if os.path.isfile(executable)  and  (not execute('Getting python configuration info...', '{executable}-config --prefix --includes'.format(**vars()), terminate_on_failure=False) ):
+    #     print('found executable!')
+    #     _, executable_version = execute('Checking Python interpreter version...', '{executable} --version'.format(**vars()), return_='tuple')
+    #     executable_version = executable_version.split()[-1]
+    # else: executable_version = ''
+    # print('executable_version: {}'.format(executable_version))
+    #if executable_version != url.rpartition('Python-')[2][:-len('.tgz')]:
+
+    if os.path.isfile(signature_file_name) and open(signature_file_name).read() == signature:
+        #print('Install for Python-{} is detected, skipping installation procedure...'.format(python_version))
+        pass
+
+    else:
+        print( 'Installing Python-{python_version}, using {url} with extra:{extra}...'.format( **vars() ) )
+
+        if os.path.isdir(root): shutil.rmtree(root)
+
+        build_prefix = os.path.abspath(root + '/../build-python-{}'.format(python_version) )
+
+        if not os.path.isdir(root): os.makedirs(root)
+        if not os.path.isdir(build_prefix): os.makedirs(build_prefix)
+
+        platform_is_mac = True if platform['os'] in ['mac', 'm1'] else False
+        platform_is_linux = not platform_is_mac
+
+        #if False and platform['os'] == 'mac' and platform_module.machine() == 'arm64' and tuple( map(int, python_version.split('.') ) ) >= (3, 9):
+        if ( platform['os'] == 'mac' and python_version == '3.6' ) \
+           or ( platform_is_linux and python_version in ['3.10', '3.11'] ):
+            open_ssl_url = local_open_ssl_install(root, build_prefix, jobs)
+            options += f' --with-openssl={root} --with-openssl-rpath=auto'
+            #signature += 'OpenSSL install: ' + open_ssl_url + '\n'
+
+        archive = build_prefix + '/' + url.split('/')[-1]
+        build_dir = archive.rpartition('.tgz')[0]
+        if os.path.isdir(build_dir): shutil.rmtree(build_dir)
+
+        with open(archive, 'wb') as f:
+            #response = urllib2.urlopen(url)
+            response = urllib.request.urlopen(url)
+            f.write( response.read() )
+
+        #execute('Execution environment:', 'env'.format(**vars()) )
+
+        execute('Unpacking {}'.format(archive), 'cd {build_prefix} && tar -xvzf {archive}'.format(**vars()) )
+
+        #execute('Building and installing...', 'cd {} && CC={compiler} CXX={cpp_compiler} {extra[0]} ./configure {extra[1]} --prefix={root} && {extra[0]} make -j{jobs} && {extra[0]} make install'.format(build_dir, **locals()) )
+        execute('Configuring...', 'cd {} && CC={compiler} CXX={cpp_compiler} {extra[0]} ./configure {options} {extra[1]} --prefix={root}'.format(build_dir, **locals()) )
+        execute('Building...', 'cd {} && {extra[0]} make -j{jobs}'.format(build_dir, **locals()) )
+        execute('Installing...', 'cd {} && {extra[0]} make -j{jobs} install'.format(build_dir, **locals()) )
+
+        shutil.rmtree(build_prefix)
+
+        #execute('Updating setuptools...', f'cd {root} && {root}/bin/pip{python_version} install --upgrade setuptools wheel' )
+
+        # if 'certifi' not in packages:
+        #     packages += ' certifi'
+
+        if packages: execute( f'Installing packages {packages}...', f'cd {root} && unset __PYVENV_LAUNCHER__ && {root}/bin/pip{python_version} install --upgrade {packages}' )
+        #if packages: execute( f'Installing packages {packages}...', f'cd {root} && unset __PYVENV_LAUNCHER__ && {executable} -m pip install --upgrade {packages}' )
+
+        remove_pip_and_easy_install(root)  # removing all pip's and easy_install's to make sure that environment is immutable
+
+        with open(signature_file_name, 'w') as f: f.write(signature)
+
+        print( 'Installing Python-{python_version}, using {url} with extra:{extra}... Done.'.format( **vars() ) )
+
+    il = get_python_include_and_lib(executable)
+
+    return NT(
+        python = executable,
+        root = root,
+        python_include_dir = il.python_include_dir,
+        python_lib_dir = il.python_lib_dir,
+        version = python_version,
+        url = url,
+        platform = platform,
+        config = config,
+        hash = hash,
+    )
+
+
+
+def setup_python_virtual_environment(working_dir, python_environment, packages=''):
+    ''' Deploy Python virtual environment at working_dir
+    '''
+
+    python = python_environment.python
+
+    execute('Setting up Python virtual environment...', 'unset __PYVENV_LAUNCHER__ && {python} -m venv --clear {working_dir}'.format(**vars()) )
+
+    activate = f'unset __PYVENV_LAUNCHER__ && . {working_dir}/bin/activate'
+
+    bin=working_dir+'/bin'
+
+    if packages: execute('Installing packages: {}...'.format(packages), 'unset __PYVENV_LAUNCHER__ && {bin}/python {bin}/pip install --upgrade pip setuptools && {bin}/python {bin}/pip install --progress-bar off {packages}'.format(**vars()) )
+    #if packages: execute('Installing packages: {}...'.format(packages), '{bin}/pip{python_environment.version} install {packages}'.format(**vars()) )
+
+    return NT(activate = activate, python = bin + '/python', root = working_dir, bin = bin)
+
+
+
+def setup_persistent_python_virtual_environment(python_environment, packages):
+    ''' Setup _persistent_ and _immutable_ Python virtual environment which will be saved between test runs
+    '''
+
+    if python_environment.version.startswith('2.'):
+        assert not packages, f'ERROR: setup_persistent_python_virtual_environment does not support Python-2.* with non-empty package list!'
+        return NT(activate = ':', python = python_environment.python, root = python_environment.root, bin = python_environment.root + '/bin')
+
+    else:
+        #if 'certifi' not in packages: packages += ' certifi'
+
+        h = hashlib.md5()
+        h.update(f'v1.0.0 platform: {python_environment.platform} python_source_url: {python_environment.url} python-hash: {python_environment.hash} packages: {packages}'.encode('utf-8', errors='backslashreplace') )
+        hash = h.hexdigest()
+
+        prefix = calculate_unique_prefix_path(python_environment.platform, python_environment.config)
+
+        root = os.path.abspath( prefix + '/python_virtual_environments/' + '/python-' + python_environment.version + '/' + hash )
+        signature_file_name = root + '/.signature'
+        signature = f'setup_persistent_python_virtual_environment v1.0.0\npython: {python_environment.hash}\npackages: {packages}\n'
+
+        activate = f'unset __PYVENV_LAUNCHER__ && . {root}/bin/activate'
+        bin = f'{root}/bin'
+
+        if os.path.isfile(signature_file_name) and open(signature_file_name).read() == signature: pass
+        else:
+            if os.path.isdir(root): shutil.rmtree(root)
+            setup_python_virtual_environment(root, python_environment, packages=packages)
+            remove_pip_and_easy_install(root)  # removing all pip's and easy_install's to make sure that environment is immutable
+            with open(signature_file_name, 'w') as f: f.write(signature)
+
+        return NT(activate = activate, python = bin + '/python', root = root, bin = bin, hash = hash)
+
+
+
+def _get_path_to_conda_root(platform, config):
+    ''' Perform local (prefix) install of miniconda and return NT(activate, conda_root_dir, conda)
+        this function is for inner use only, - to setup custom conda environment inside your test use `setup_conda_virtual_environment` defined below
+    '''
+    miniconda_sources = {
+        'mac'    : 'https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh',
+        'linux'  : 'https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh',
+        'aarch64': 'https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh',
+        'ubuntu' : 'https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh',
+        'm1'     : 'https://repo.anaconda.com/miniconda/Miniconda3-py38_4.10.1-MacOSX-arm64.sh',
+    }
+
+    conda_sources = {
+        'mac'    : 'https://repo.continuum.io/archive/Anaconda3-2018.12-MacOSX-x86_64.sh',
+        'linux'  : 'https://repo.continuum.io/archive/Anaconda3-2018.12-Linux-x86_64.sh',
+        'ubuntu' : 'https://repo.continuum.io/archive/Anaconda3-2018.12-Linux-x86_64.sh',
+    }
+
+    #platform_os = 'm1' if platform_module.machine() == 'arm64' else platform['os']
+    #url = miniconda_sources[ platform_os ]
+
+    platform_os = platform['os']
+    for o in 'alpine centos ubuntu'.split():
+        if platform_os.startswith(o): platform_os = 'linux'
+
+    url = miniconda_sources[platform_os]
+
+    version = '1'
+    channels = ''  # conda-forge
+
+    #packages = ['conda-build gcc libgcc', 'libgcc=5.2.0'] # libgcc installs is workaround for "Anaconda libstdc++.so.6: version `GLIBCXX_3.4.20' not found", see: https://stackoverflow.com/questions/48453497/anaconda-libstdc-so-6-version-glibcxx-3-4-20-not-found
+    #packages = ['conda-build gcc'] # libgcc installs is workaround for "Anaconda libstdc++.so.6: version `GLIBCXX_3.4.20' not found", see: https://stackoverflow.com/questions/48453497/anaconda-libstdc-so-6-version-glibcxx-3-4-20-not-found
+    packages = ['conda-build anaconda-client conda-verify',]
+
+    signature = f'url: {url}\nversion: {version}\channels: {channels}\npackages: {packages}\n'
+
+    root = calculate_unique_prefix_path(platform, config) + '/conda'
+
+    signature_file_name = root + '/.signature'
+
+    # presense of __PYVENV_LAUNCHER__,PYTHONHOME, PYTHONPATH sometimes confuse Python so we have to unset them
+    unset = 'unset __PYVENV_LAUNCHER__ && unset PYTHONHOME && unset PYTHONPATH'
+    activate = unset + ' && . ' + root + '/bin/activate'
+
+    executable = root + '/bin/conda'
+
+
+    if os.path.isfile(signature_file_name) and open(signature_file_name).read() == signature:
+        print( f'Install for MiniConda is detected, skipping installation procedure...' )
+
+    else:
+        print( f'Installing MiniConda, using {url}...' )
+
+        if os.path.isdir(root): shutil.rmtree(root)
+
+        build_prefix = os.path.abspath(root + f'/../build-conda' )
+
+        #if not os.path.isdir(root): os.makedirs(root)
+        if not os.path.isdir(build_prefix): os.makedirs(build_prefix)
+
+        archive = build_prefix + '/' + url.split('/')[-1]
+
+        with open(archive, 'wb') as f:
+            response = urllib.request.urlopen(url)
+            f.write( response.read() )
+
+        execute('Installing conda...', f'cd {build_prefix} && {unset} && bash {archive} -b -p {root}' )
+
+        # conda update --yes --quiet -n base -c defaults conda
+
+        if channels: execute(f'Adding extra channles {channels}...', f'cd {build_prefix} && {activate} && conda config --add channels {channels}' )
+
+        for p in packages: execute(f'Installing conda packages: {p}...', f'cd {build_prefix} && {activate} && conda install --quiet --yes {p}' )
+
+        shutil.rmtree(build_prefix)
+
+        with open(signature_file_name, 'w') as f: f.write(signature)
+
+        print( f'Installing MiniConda, using {url}... Done.' )
+
+    execute(f'Updating conda base...', f'{activate} && conda update --all --yes' )
+    return NT(conda=executable, root=root, activate=activate, url=url)
+
+
+
+def setup_conda_virtual_environment(working_dir, platform, config, packages=''):
+    ''' Deploy Conda virtual environment at working_dir
+    '''
+    conda_root_env = _get_path_to_conda_root(platform, config)
+    activate = conda_root_env.activate
+
+    python_version = platform.get('python', DEFAULT_PYTHON_VERSION)
+
+    prefix = os.path.abspath( working_dir + '/.conda-python-' + python_version )
+
+    command_line = f'conda create --quiet --yes --prefix {prefix} python={python_version}'
+
+    execute( f'Setting up Conda for Python-{python_version} virtual environment...', f'cd {working_dir} && {activate} && ( {command_line} || ( conda clean --yes && {command_line} ) )' )
+
+    activate = f'{activate} && conda activate {prefix}'
+
+    if packages: execute( f'Setting up extra packages {packages}...', f'cd {working_dir} && {activate} && conda install --quiet --yes {packages}' )
+
+    python = prefix + '/bin/python' + python_version
+
+    il = get_python_include_and_lib(python)
+
+    return NT(
+        activate = activate,
+        root = prefix,
+        python = python,
+        python_include_dir = il.python_include_dir,
+        python_lib_dir = il.python_lib_dir,
+        version = python_version,
+        activate_base = conda_root_env.activate,
+        url = prefix, # conda_root_env.url,
+        platform=platform,
+        config=config,
+    )
+
+
+
+class FileLock():
+    ''' Implementation of file-lock object that could be use with Python `with` statement
+    '''
+
+    def __init__(self, file_name):
+        self.locked = False
+        self.file_name = file_name
+
+
+    def __enter__(self):
+        if not self.locked: self.acquire()
+        return self
+
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.locked: self.release()
+
+
+    def __del__(self):
+        self.release()
+
+
+    def acquire(self):
+        while True:
+            try:
+                os.close( os.open(self.file_name, os.O_CREAT | os.O_EXCL, mode=0o600) )
+                self.locked = True
+                break
+
+            except FileExistsError as e:
+                time.sleep(60)
+
+
+    def release(self):
+        if self.locked:
+            os.remove(self.file_name)
+            self.locked = False
+
+
+
+def convert_submodule_urls_from_ssh_to_https(repository_root):
+    ''' switching submodules URL to HTTPS so we can clone without SSH key
+    '''
+    with open(f'{repository_root}/.gitmodules') as f: m = f.read()
+    with open(f'{repository_root}/.gitmodules', 'w') as f:
+        f.write(
+            m
+            .replace('url = git@github.com:', 'url = https://github.com/')
+            .replace('url = ../../../',       'url = https://github.com/RosettaCommons/')
+            .replace('url = ../../',          'url = https://github.com/RosettaCommons/')
+            .replace('url = ../',             'url = https://github.com/RosettaCommons/')
+        )

.rosetta-ci/tests/rfd.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+# (c) Copyright Rosetta Commons Member Institutions.
+# (c) This file is part of the Rosetta software suite and is made available under license.
+# (c) The Rosetta software is developed by the contributing members of the Rosetta Commons.
+# (c) For more information, see http://www.rosettacommons.org. Questions about this can be
+# (c) addressed to University of Washington CoMotion, email: license@uw.edu.
+
+## @file   rfd.py
+## @brief  main test files for RFdiffusion
+## @author Sergey Lyskov
+
+
+import imp
+imp.load_source(__name__, '/'.join(__file__.split('/')[:-1]) +  '/__init__.py')  # A bit of Python magic here, what we trying to say is this: from __init__ import *, but init is calculated from file location
+
+_api_version_ = '1.0'
+
+import os, tempfile, shutil
+import urllib.request
+
+
+_models_urls_ = '''
+http://files.ipd.uw.edu/pub/RFdiffusion/6f5902ac237024bdd0c176cb93063dc4/Base_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/e29311f6f1bf1af907f9ef9f44b8328b/Complex_base_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/60f09a193fb5e5ccdc4980417708dbab/Complex_Fold_base_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/74f51cfb8b440f50d70878e05361d8f0/InpaintSeq_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/76d00716416567174cdb7ca96e208296/InpaintSeq_Fold_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/5532d2e1f3a4738decd58b19d633b3c3/ActiveSite_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/12fc204edeae5b57713c5ad7dcb97d39/Base_epoch8_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/f572d396fae9206628714fb2ce00f72e/Complex_beta_ckpt.pt
+http://files.ipd.uw.edu/pub/RFdiffusion/1befcb9b28e2f778f53d47f18b7597fa/RF_structure_prediction_weights.pt
+'''.split()
+
+
+def run_main_test_suite(repository_root, working_dir, platform, config, debug):
+    full_log = ''
+
+    python_environment = local_python_install(platform, config)
+
+    models_dir = repository_root + '/models'
+    if not os.path.isdir(models_dir): os.makedirs(models_dir)
+
+    for url in _models_urls_:
+        file_name = models_dir + '/' + url.split('/')[-1]
+        tmp_file_name = file_name + '.tmp'
+        if not os.path.isfile(file_name):
+            print(f'downloading {url}...')
+            full_log += f'downloading {url}...\n'
+            urllib.request.urlretrieve(url, tmp_file_name)
+            os.rename(tmp_file_name, file_name)
+
+    execute('unpacking ppi scaffolds...', f'cd {repository_root} && tar -xvf examples/ppi_scaffolds_subset.tar.gz -C examples')
+
+    with tempfile.TemporaryDirectory(dir=working_dir) as tmpdirname:
+    # tmpdirname = working_dir+'/.ve'
+    # if True:
+
+        #ve = setup_persistent_python_virtual_environment(python_environment, packages='numpy torch omegaconf scipy opt_einsum dgl')
+        #ve = setup_python_virtual_environment(working_dir+'/.ve', python_environment, packages='numpy torch omegaconf scipy opt_einsum dgl e3nn icecream pyrsistent wandb pynvml decorator jedi hydra-core')
+        ve = setup_python_virtual_environment(tmpdirname, python_environment, packages='numpy torch omegaconf scipy opt_einsum dgl e3nn icecream pyrsistent wandb pynvml decorator jedi hydra-core')
+
+        execute('Installing local se3-transformer package...', f'cd {repository_root}/env/SE3Transformer && {ve.bin}/pip3 install --editable .')
+        execute('Installing RFdiffusion package...', f'cd {repository_root} && {ve.bin}/pip3 install --editable .')
+
+        #res, output = execute('running unit tests...', f'{ve.activate} && cd {repository_root} && python -m unittest', return_='tuple', add_message_and_command_line_to_output=True)
+        #res, output = execute('running unit tests...', f'cd {repository_root} && {ve.bin}/pytest', return_='tuple')
+
+
+        results_file = f'{repository_root}/tests/.results.json'
+        if os.path.isfile(results_file): os.remove(results_file)
+
+        res, output = execute('running RFdiffusion tests...', f'{ve.activate} && cd {repository_root}/tests && python test_diffusion.py', return_='tuple', add_message_and_command_line_to_output=True)
+
+        if os.path.isfile(results_file):
+            with open(results_file) as f: sub_tests_reults = json.load(f)
+
+            state = _S_passed_
+            for r in sub_tests_reults.values():
+                if r[_StateKey_] == _S_failed_:
+                    state = _S_failed_
+                    break
+
+        else:
+            sub_tests_reults = {}
+            output  += '\n\nEmpty sub-test results, marking test as `failed`...'
+            state = _S_failed_
+
+        shutil.move(f'{repository_root}/tests/outputs', f'{working_dir}/outputs')
+
+        for d in os.listdir(f'{repository_root}/tests'):
+            p = f'{repository_root}/tests/{d}'
+            if d.startswith('tests_') and os.path.isdir(p): shutil.rmtree(p)
+
+        results = {
+            _StateKey_ : state,
+            _LogKey_ : full_log + '\n' + output,
+            _ResultsKey_ : {
+                _TestsKey_ : sub_tests_reults,
+            },
+        }
+
+    return results
+
+
+
+def run(test, repository_root, working_dir, platform, config, hpc_driver=None, verbose=False, debug=False):
+    if test == '': return run_main_test_suite(repository_root=repository_root, working_dir=working_dir, platform=platform, config=config, debug=debug)
+    else: raise BenchmarkError('Unknow scripts test: {}!'.format(test))

.rosetta-ci/tests/self.md

@@ -0,0 +1,6 @@
+# self test suite
+These tests are design to help debug interface between testing server and Rosetta testing scripts
+
+-----
+### python
+Test Python platform support and functionality of local and persistent Python virtual environments

.rosetta-ci/tests/self.py

@@ -0,0 +1,209 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# :noTabs=true:
+
+# (c) Copyright Rosetta Commons Member Institutions.
+# (c) This file is part of the Rosetta software suite and is made available under license.
+# (c) The Rosetta software is developed by the contributing members of the Rosetta Commons.
+# (c) For more information, see http://www.rosettacommons.org. Questions about this can be
+# (c) addressed to University of Washington CoMotion, email: license@uw.edu.
+
+## @file   dummy.py
+## @brief  self-test and debug-aids tests
+## @author Sergey Lyskov
+
+import os, os.path, shutil, re, string
+import json
+
+import random
+
+import imp
+imp.load_source(__name__, '/'.join(__file__.split('/')[:-1]) +  '/__init__.py')  # A bit of Python magic here, what we trying to say is this: from __init__ import *, but init is calculated from file location
+
+_api_version_ = '1.0'
+
+
+def run_state_test(repository_root, working_dir, platform, config):
+    revision_id = config['revision']
+    states = (_S_passed_, _S_failed_, _S_build_failed_, _S_script_failed_)
+    state = states[revision_id % len(states)]
+
+    return {_StateKey_ : state,  _ResultsKey_ : {},  _LogKey_ : f'run_state_test: setting test state to {state!r}...' }
+
+
+sub_test_description_template = '''\
+# subtests_test test suite
+These sub-test description is generated for 3/4 of sub-tests
+
+-----
+### {name}
+The warm time, had already disappeared like dust. Broken rain, fragment of light shadow, bring more pain to my heart...
+-----
+'''
+
+def run_subtests_test(repository_root, working_dir, platform, config):
+    tests = {}
+    for i in range(16):
+        name = f's-{i:02}'
+        log = ('x'*63 + '\n') * 16 * 256 * i
+        s = i % 3
+        if   s == 0: state = _S_passed_
+        elif s == 1: state = _S_failed_
+        else:        state = _S_script_failed_
+
+        if i % 4:
+            os.mkdir( f'{working_dir}/{name}' )
+            with open(f'{working_dir}/{name}/description.md', 'w') as f: f.write( sub_test_description_template.format(**vars()) )
+
+            with open( f'{working_dir}/{name}/fantome.txt', 'w') as f: f.write('No one wants to hear the sequel to a fairytale\n')
+
+        tests[name] = { _StateKey_ : state, _LogKey_   : log, }
+
+    test_log = ('*'*63 + '\n') * 16 * 1024 * 16
+    return {_StateKey_ : _S_failed_,  _ResultsKey_ : {_TestsKey_: tests},  _LogKey_ : test_log }
+
+
+def run_regression_test(repository_root, working_dir, platform, config):
+    const     = 'const'
+    volatile  = 'volatile'
+    new       = ''.join( random.sample( string.ascii_letters + string.digits, 8) )
+    oversized = 'oversized'
+
+    sub_tests = [const, volatile, new]
+
+    const_dir = working_dir + '/' + const
+    os.mkdir(const_dir)
+    with open(const_dir + '/const_data', 'w') as f: f.write( '\n'.join( (str(i) for i in range(32) ) ) )
+
+    volatile_dir = working_dir + '/' + volatile
+    os.mkdir(volatile_dir)
+    with open(volatile_dir + '/const_data', 'w') as f: f.write( '\n'.join( (str(i) for i in range(32, 64) ) ) )
+    with open(volatile_dir + '/volatile_data', 'w') as f: f.write( '\n'.join( ( ''.join(random.sample( string.ascii_letters + string.digits, 8) ) for i in range(32) ) ) )
+
+    new_dir = working_dir + '/' + new
+    os.mkdir(new_dir)
+    with open(new_dir + '/data', 'w') as f: f.write( '\n'.join( (str(i) for i in range(64)) ) )
+
+
+    new_dir = working_dir + '/' + oversized
+    os.mkdir(new_dir)
+    with open(new_dir + '/large', 'w') as f: f.write( ('x'*63 + '\n')*16*1024*256 +'extra')
+
+    return {_StateKey_ : _S_queued_for_comparison_,  _ResultsKey_ : {},  _LogKey_ : f'sub-tests: {sub_tests!r}' }
+
+
+
+def run_release_test(repository_root, working_dir, platform, config):
+    release_root = config['mounts'].get('release_root')
+
+    branch = config['branch']
+    revision = config['revision']
+
+    assert release_root, "config['release_root'] must be set!"
+
+    release_path = f'{release_root}/dummy'
+
+    if not os.path.isdir(release_path): os.makedirs(release_path)
+
+    with open(f'{release_path}/{branch}-{revision}.txt', 'w') as f: f.write('dummy release file\n')
+
+    return {_StateKey_ : _S_passed_,  _ResultsKey_ : {},  _LogKey_ : f'Config release root set to: {release_root}'}
+
+
+
+def run_python_test(repository_root, working_dir, platform, config):
+
+    import zlib, ssl
+
+    python_environment = local_python_install(platform, config)
+
+    if platform['python'][0] == '2': pass
+    else:
+
+        if platform['os'] == 'mac' and int( platform['python'].split('.')[1] ) > 6 :
+            # SSL certificate test
+            import urllib.request; urllib.request.urlopen('https://benchmark.graylab.jhu.edu')
+
+        ves = [
+            setup_persistent_python_virtual_environment(python_environment, packages='colr dice xdice pdp11games'),
+            setup_python_virtual_environment(working_dir, python_environment, packages='colr dice xdice pdp11games'),
+        ]
+
+        for ve in ves:
+            commands = [
+                'import colr, dice, xdice, pdp11games',
+            ]
+
+            if platform['os'] == 'mac' and int( platform['python'].split('.')[1] ) > 6 :
+                # SSL certificate test
+                commands.append('import urllib.request; urllib.request.urlopen("https://benchmark.graylab.jhu.edu/queue")')
+
+            for command in commands:
+                execute('Testing local Python virtual enviroment...', f"{ve.activate} && {ve.python} -c '{command}'")
+                execute('Testing local Python virtual enviroment...', f"{ve.activate} && python -c '{command}'")
+
+
+
+    return {_StateKey_ : _S_passed_,  _ResultsKey_ : {},  _LogKey_ : f'Done!'}
+
+
+
+def compare(test, results, files_path, previous_results, previous_files_path):
+    """
+    Compare the results of two tests run (new vs. previous) for regression test
+    Take two dict and two paths
+    Must return standard dict with results
+
+    :param test: str
+    :param results: dict
+    :param files_path: str
+    :param previous_results: dict
+    :param previous_files_path: str
+    :rtype: dict
+    """
+    ignore_files = []
+
+    results = dict(tests={}, summary=dict(total=0, failed=0, failed_tests=[]))  # , config={}
+
+    if previous_files_path:
+        for test in os.listdir(files_path):
+            if os.path.isdir(files_path + '/' + test):
+                exclude = ''.join([' --exclude="{}"'.format(f) for f in ignore_files] ) + ' --exclude="*.ignore"'
+                res, brief_diff = execute('Comparing {}...'.format(test), 'diff -rq {exclude} {0}/{test} {1}/{test}'.format(previous_files_path, files_path, test=test, exclude=exclude), return_='tuple')
+                res, full_diff  = execute('Comparing {}...'.format(test), 'diff -r  {exclude} {0}/{test} {1}/{test}'.format(previous_files_path, files_path, test=test, exclude=exclude), return_='tuple')
+                diff = 'Brief Diff:\n' + brief_diff + ( ('\n\nFull Diff:\n' + full_diff[:1024*1024*1]) if full_diff != brief_diff else '' )
+
+                state = _S_failed_ if res else _S_passed_
+                results['tests'][test] = {_StateKey_: state, _LogKey_: diff if state != _S_passed_ else ''}
+
+                results['summary']['total'] += 1
+                if res: results['summary']['failed'] += 1; results['summary']['failed_tests'].append(test)
+
+    else: # no previous tests case, returning 'passed' for all sub_tests
+        for test in os.listdir(files_path):
+            if os.path.isdir(files_path + '/' + test):
+                results['tests'][test] = {_StateKey_: _S_passed_, _LogKey_: 'First run, no previous results available. Skipping comparison...\n'}
+                results['summary']['total'] += 1
+
+    for test in os.listdir(files_path):
+        if os.path.isdir(files_path + '/' + test):
+            if os.path.isfile(files_path+'/'+test+'/.test_did_not_run.log')  or  os.path.isfile(files_path+'/'+test+'/.test_got_timeout_kill.log'):
+                results['tests'][test][_StateKey_] = _S_script_failed_
+                results['tests'][test][_LogKey_] += '\nCompare(...): Marking as "Script failed" due to presense of .test_did_not_run.log or .test_got_timeout_kill.log file!\n'
+                if test not in results['summary']['failed_tests']:
+                    results['summary']['failed'] += 1
+                    results['summary']['failed_tests'].append(test)
+
+    state = _S_failed_ if results['summary']['failed'] else _S_passed_
+
+    return {_StateKey_: state, _LogKey_: 'Comparison dummy log...', _ResultsKey_: results}
+
+
+def run(test, repository_root, working_dir, platform, config, hpc_driver=None, verbose=False, debug=False):
+    if   test == 'state':      return run_state_test      (repository_root, working_dir, platform, config)
+    elif test == 'regression': return run_regression_test (repository_root, working_dir, platform, config)
+    elif test == 'subtests':   return run_subtests_test   (repository_root, working_dir, platform, config)
+    elif test == 'release':    return run_release_test    (repository_root, working_dir, platform, config)
+    elif test == 'python':     return run_python_test     (repository_root, working_dir, platform, config)
+
+    else: raise BenchmarkError(f'Dummy test script does not support run with test={test!r}!')

END

@@ -0,0 +1,7 @@
+{
+	"retCode":"100",
+	"retData":null,
+	"retMsg":"操作成功",
+	"retTime":"2022-11-05 22:20:09",
+	"success":true
+}

LICENSE

@@ -0,0 +1,30 @@
+BSD License
+
+Copyright (c) 2023 University of Washington. Developed at the Institute for
+Protein Design by Joseph Watson, David Juergens, Nathaniel Bennett, Brian Trippe
+and Jason Yim
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+Redistributions of source code must retain the above copyright notice, this
+list of conditions and the following disclaimer.
+
+Redistributions in binary form must reproduce the above copyright notice, this
+list of conditions and the following disclaimer in the documentation and/or
+other materials provided with the distribution.
+
+Neither the name of the University of Washington nor the names of its
+contributors may be used to endorse or promote products derived from this
+software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE UNIVERSITY OF WASHINGTON AND CONTRIBUTORS “AS
+IS” AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE UNIVERSITY OF WASHINGTON OR CONTRIBUTORS BE
+LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE
+GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
+HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
+OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

README.md

@@ -1,3 +1,516 @@
+# RF*diffusion*
+
+<!--
+<img width="1115" alt="Screen Shot 2023-01-19 at 5 56 33 PM" src="https://user-images.githubusercontent.com/56419265/213588200-f8f44dba-276e-4dd2-b844-15acc441458d.png">
+-->
+<p align="center">
+  <img src="./img/diffusion_protein_gradient_2.jpg" alt="alt text" width="1100px" align="middle"/>
+</p>
+
+*Image: Ian C. Haydon / UW Institute for Protein Design*
+
+## Description
+
+RFdiffusion is an open source method for structure generation, with or without conditional information (a motif, target etc). It can perform a whole range of protein design challenges as we have outlined in [the RFdiffusion paper](https://www.biorxiv.org/content/10.1101/2022.12.09.519842v1).
+
+**Things Diffusion can do**
+- Motif Scaffolding
+- Unconditional protein generation
+- Symmetric unconditional generation (cyclic, dihedral and tetrahedral symmetries currently implemented, more coming!)
+- Symmetric motif scaffolding
+- Binder design
+- Design diversification ("partial diffusion", sampling around a design)
+
+----
+
+# Table of contents
+
+- [RF*diffusion*](#rfdiffusion)
+  - [Description](#description)
+- [Table of contents](#table-of-contents)
+- [Getting started / installation](#getting-started--installation)
+    - [Conda Install SE3-Transformer](#conda-install-se3-transformer)
+    - [Get PPI Scaffold Examples](#get-ppi-scaffold-examples)
+- [Usage](#usage)
+    - [Running the diffusion script](#running-the-diffusion-script)
+    - [Basic execution - an unconditional monomer](#basic-execution---an-unconditional-monomer)
+    - [Motif Scaffolding](#motif-scaffolding)
+    - [The "active site" model holds very small motifs in place](#the-active-site-model-holds-very-small-motifs-in-place)
+    - [The `inpaint_seq` flag](#the-inpaint_seq-flag)
+    - [A note on `diffuser.T`](#a-note-on-diffusert)
+    - [Partial diffusion](#partial-diffusion)
+    - [Binder Design](#binder-design)
+    - [Practical Considerations for Binder Design](#practical-considerations-for-binder-design)
+    - [Fold Conditioning](#fold-conditioning)
+    - [Generation of Symmetric Oligomers](#generation-of-symmetric-oligomers)
+    - [Using Auxiliary Potentials](#using-auxiliary-potentials)
+    - [Symmetric Motif Scaffolding.](#symmetric-motif-scaffolding)
+    - [A Note on Model Weights](#a-note-on-model-weights)
+    - [Things you might want to play with at inference time](#things-you-might-want-to-play-with-at-inference-time)
+    - [Understanding the output files](#understanding-the-output-files)
+    - [Docker](#docker)
+    - [Conclusion](#conclusion)
+
+# Getting started / installation
+
+Thanks to Sergey Ovchinnikov, RFdiffusion is available as a [Google Colab Notebook](https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/rf/examples/diffusion.ipynb) if you would like to run it there!
+
+We strongly recommend reading this README carefully before getting started with RFdiffusion, and working through some of the examples in the Colab Notebook.
+
+If you want to set up RFdiffusion locally, follow the steps below:
+
+To get started using RFdiffusion, clone the repo:
+```
+git clone https://github.com/RosettaCommons/RFdiffusion.git
+```
+
+You'll then need to download the model weights into the RFDiffusion directory.
+```
+cd RFdiffusion
+mkdir models && cd models
+wget http://files.ipd.uw.edu/pub/RFdiffusion/6f5902ac237024bdd0c176cb93063dc4/Base_ckpt.pt
+wget http://files.ipd.uw.edu/pub/RFdiffusion/e29311f6f1bf1af907f9ef9f44b8328b/Complex_base_ckpt.pt
+wget http://files.ipd.uw.edu/pub/RFdiffusion/60f09a193fb5e5ccdc4980417708dbab/Complex_Fold_base_ckpt.pt
+wget http://files.ipd.uw.edu/pub/RFdiffusion/74f51cfb8b440f50d70878e05361d8f0/InpaintSeq_ckpt.pt
+wget http://files.ipd.uw.edu/pub/RFdiffusion/76d00716416567174cdb7ca96e208296/InpaintSeq_Fold_ckpt.pt
+wget http://files.ipd.uw.edu/pub/RFdiffusion/5532d2e1f3a4738decd58b19d633b3c3/ActiveSite_ckpt.pt
+wget http://files.ipd.uw.edu/pub/RFdiffusion/12fc204edeae5b57713c5ad7dcb97d39/Base_epoch8_ckpt.pt
+
+Optional:
+wget http://files.ipd.uw.edu/pub/RFdiffusion/f572d396fae9206628714fb2ce00f72e/Complex_beta_ckpt.pt
+
+# original structure prediction weights
+wget http://files.ipd.uw.edu/pub/RFdiffusion/1befcb9b28e2f778f53d47f18b7597fa/RF_structure_prediction_weights.pt
+```
+
+
+### Conda Install SE3-Transformer
+
+Ensure that you have either [Anaconda or Miniconda](https://conda.io/projects/conda/en/latest/user-guide/install/index.html) installed.
+
+You also need to install [NVIDIA's implementation of SE(3)-Transformers](https://developer.nvidia.com/blog/accelerating-se3-transformers-training-using-an-nvidia-open-source-model-implementation/) Here is how to install the NVIDIA SE(3)-Transformer code:
+
+```
+conda env create -f env/SE3nv.yml
+
+conda activate SE3nv
+cd env/SE3Transformer
+pip install --no-cache-dir -r requirements.txt
+python setup.py install
+cd ../.. # change into the root directory of the repository
+pip install -e . # install the rfdiffusion module from the root of the repository
+```
+Anytime you run diffusion you should be sure to activate this conda environment by running the following command:
+```
+conda activate SE3nv
+```
+Total setup should take less than 30 minutes on a standard desktop computer.
+Note: Due to the variation in GPU types and drivers that users have access to, we are not able to make one environment that will run on all setups. As such, we are only providing a yml file with support for CUDA 11.1 and leaving it to each user to customize it to work on their setups. This customization will involve changing the cudatoolkit and (possibly) the PyTorch version specified in the yml file.
+
 ---
-license: afl-3.0
+
+### Get PPI Scaffold Examples
+
+To run the scaffolded protein binder design (PPI) examples, we have provided some example scaffold files (`examples/ppi_scaffolds_subset.tar.gz`).
+You'll need to untar this:
+```
+tar -xvf examples/ppi_scaffolds_subset.tar.gz -C examples/
+```
+
+We will explain what these files are and how to use them in the Fold Conditioning section.
+
+----
+
+
+# Usage
+In this section we will demonstrate how to run diffusion.
+
+
+
+<p align="center">
+  <img src="./img/main.png" alt="alt text" width="1100px" align="middle"/>
+</p>
+
+
+### Running the diffusion script
+The actual script you will execute is called `scripts/run_inference.py`. There are many ways to run it, governed by hydra configs.
+[Hydra configs](https://hydra.cc/docs/configure_hydra/intro/) are a nice way of being able to specify many different options, with sensible defaults drawn *directly* from the model checkpoint, so inference should always, by default, match training.
+What this means is that the default values in `config/inference/base.yml` might not match the actual values used during inference, with a specific checkpoint. This is all handled under the hood.
+
 ---
+### Basic execution - an unconditional monomer
+<img src="./img/cropped_uncond.png" alt="alt text" width="400px" align="right"/>
+
+Let's first look at how you would do unconditional design of a protein of length 150aa.
+For this, we just need to specify three things:
+1. The length of the protein
+2. The location where we want to write files to
+3. The number of designs we want
+
+```
+./scripts/run_inference.py 'contigmap.contigs=[150-150]' inference.output_prefix=test_outputs/test inference.num_designs=10
+```
+
+Let's look at this in detail.
+Firstly, what is `contigmap.contigs`?
+Hydra configs tell the inference script how it should be run. To keep things organised, the config has different sub-configs, one of them being `contigmap`, which pertains to everything related to the contig string (that defines the protein being built).
+Take a look at the config file if this isn't clear: `configs/inference/base.yml`
+Anything in the config can be overwritten manually from the command line. You could, for example, change how the diffuser works:
+```
+diffuser.crd_scale=0.5
+```
+... but don't do this unless you really know what you're doing!!
+
+
+Now, what does `'contigmap.contigs=[150-150]'` mean?
+To those who have used RFjoint inpainting, this might look familiar, but a little bit different. Diffusion, in fact, uses the identical 'contig mapper' as inpainting, except that, because we're using hydra, we have to give this to the model in a different way. The contig string has to be passed as a single-item in a list, rather than as a string, for hydra reasons and the entire argument MUST be enclosed in `''` so that the commandline does not attempt to parse any of the special characters.
+
+The contig string allows you to specify a length range, but here, we just want a protein of 150aa in length, so you just specify [150-150]
+This will then run 10 diffusion trajectories, saving the outputs to your specified output folder.
+
+NB the first time you run RFdiffusion, it will take a while 'Calculating IGSO3'. Once it has done this, it'll be cached for future reference though! For an additional example of unconditional monomer generation, take a look at `./examples/design_unconditional.sh` in the repo! 
+
+---
+### Motif Scaffolding
+<!--
+<p align="center">
+  <img src="./img/motif.png" alt="alt text" width="700px" align="middle"/>
+</p>
+-->
+RFdiffusion can be used to scaffold motifs, in a manner akin to [Constrained Hallucination and RFjoint Inpainting](https://www.science.org/doi/10.1126/science.abn2100#:~:text=The%20binding%20and%20catalytic%20functions%20of%20proteins%20are,the%20fold%20or%20secondary%20structure%20of%20the%20scaffold.). In general, RFdiffusion significantly outperforms both Constrained Hallucination and RFjoint Inpainting.
+<p align="center">
+  <img src="./img/motif.png" alt="alt text" width="700px" align="middle"/>
+</p>
+
+When scaffolding protein motifs, we need a way of specifying that we want to scaffold some particular protein input (one or more segments from a `.pdb` file), and to be able to specify how we want these connected, and by how many residues, in the new scaffolded protein. What's more, we want to be able to sample different lengths of connecting protein, as we generally don't know *a priori* precisely how many residues we'll need to best scaffold a motif. This job of specifying inputs is handled by contigs, governed by the contigmap config in the hydra config. For those familiar with Constrained Hallucination or RFjoint Inpainting, the logic is very similar.
+Briefly:
+- Anything prefixed by a letter indicates that this is a motif, with the letter corresponding to the chain letter in the input pdb files. E.g. A10-25 pertains to residues ('A',10),('A',11)...('A',25) in the corresponding input pdb
+- Anything not prefixed by a letter indicates protein *to be built*. This can be input as a length range. These length ranges are randomly sampled each iteration of RFdiffusion inference. 
+- To specify chain breaks, we use `/0 `.
+
+In more detail, if we want to scaffold a motif, the input is just like RFjoint Inpainting, except needing to navigate the hydra config input. If we want to scaffold residues 10-25 on chain A a pdb, this would be done with `'contigmap.contigs=[5-15/A10-25/30-40]'`. This asks RFdiffusion to build 5-15 residues (randomly sampled at each inference cycle) N-terminally of A10-25 from the input pdb, followed by 30-40 residues (again, randomly sampled) to its C-terminus. If we wanted to ensure the length was always e.g. 55 residues, this can be specified with `contigmap.length=55-55`. You need to obviously also provide a path to your pdb file: `inference.input_pdb=path/to/file.pdb`. It doesn't matter if your input pdb has residues you *don't* want to scaffold - the contig map defines which residues in the pdb are actually used as the "motif". In other words, even if your pdb files has a B chain, and other residues on the A chain, *only* A10-25 will be provided to RFdiffusion.
+
+To specify that we want to inpaint in the presence of a separate chain, this can be done as follows:
+
+```
+'contigmap.contigs=[5-15/A10-25/30-40/0 B1-100]'
+```
+Look at this carefully. `/0 ` is the indicator that we want a chain break. NOTE, the space is important here. This tells the diffusion model to add a big residue jump (200aa) to the input, so that the model sees the first chain as being on a separate chain to the second.
+
+An example of motif scaffolding can be found in `./examples/design_motifscaffolding.sh`.
+
+### The "active site" model holds very small motifs in place
+In the RFdiffusion preprint we noted that for very small motifs, RFdiffusion has the tendency to not keep them perfectly fixed in the output. Therefore, for scaffolding minimalist sites such as enzyme active sites, we fine-tuned RFdiffusion on examples similar to these tasks, allowing it to hold smaller motifs better in place, and better generate *in silico* successes. If your input functional motif is very small, we reccomend using this model, which can easily be specified using the following syntax: 
+`inference.ckpt_override_path=models/ActiveSite_ckpt.pt`
+
+### The `inpaint_seq` flag
+For those familiar with RFjoint Inpainting, the contigmap.inpaint_seq input is equivalent. The idea is that often, when, for example, fusing two proteins, residues that were on the surface of a protein (and are therefore likely polar), now need to be packed into the 'core' of the protein. We therefore want them to become hydrophobic residues. What we can do, rather than directly mutating them to hydrophobics, is to mask their sequence identity, and allow RFdiffusion to implicitly reason over their sequence, and better pack against them. This requires a different model than the 'base' diffusion model, that has been trained to understand this paradigm, but this is automatically handled by the inference script (you don't need to do anything).
+
+To specify amino acids whose sequence should be hidden, use the following syntax:
+```
+'contigmap.inpaint_seq=[A1/A30-40]'
+```
+Here, we're masking the residue identity of residue A1, and all residues between A30 and A40 (inclusive). 
+
+An example of executing motif scaffolding with the `contigmap.inpaint_seq` flag is located in `./examples/design_motifscaffolding_inpaintseq.sh`
+
+### A note on `diffuser.T`
+RFdiffusion was originally trained with 200 discrete timesteps. However, recent improvements have allowed us to reduce the number of timesteps we need to use at inference time. In many cases, running with as few as approximately 20 steps provides outputs of equivalent *in silico* quality to running with 200 steps (providing a 10X speedup). The default is now set to 50 steps. Noting this is important for understanding the partial diffusion, described below. 
+
+---
+### Partial diffusion
+
+Something we can do with diffusion is to partially noise and de-noise a structure, to get some diversity around a general fold. This can work really nicely (see [Vazquez-Torres et al., BioRxiv 2022](https://www.biorxiv.org/content/10.1101/2022.12.10.519862v4.abstract)).
+This is specified by using the diffuser.parial_T input, and setting a timestep to 'noise' to. 
+<p align="center">
+  <img src="./img/partial.png" alt="alt text" width="800px" align="middle"/>
+</p>
+More noise == more diversity. In Vazquez-Torres et al., 2022, we typically used `diffuser.partial_T` of approximately 80, but this was with respect to the 200 timesteps we were using. Now that the default `diffuser.T` is 50, you will need to adjust diffuser.partial_T accordingly. E.g. now that `diffuser.T=50`, the equivalent of 80 noising steps is `diffuser.partial_T=20`. We strongly recommend sampling different values for `partial_T` however, to find the best parameters for your specific problem.
+
+When doing partial diffusion, because we are now diffusing from a known structure, this creates certain constraints. You can still use the contig input, but *this has to yield a contig string exactly the same length as the input protein*. E.g. if you have a binder:target complex, and you want to diversify the binder (length 100, chain A), you would need to input something like this:
+
+```
+'contigmap.contigs=[100-100/0 B1-150]' diffuser.partial_T=20
+```
+The reason for this is that, if your input protein was only 80 amino acids, but you've specified a desired length of 100, we don't know where to diffuse those extra 20 amino acids from, and hence, they will not lie in the distribution that RFdiffusion has learned to denoise from.
+
+An example of partial diffusion can be found in `./examples/design_partialdiffusion.sh`! 
+
+You can also keep parts of the sequence of the diffused chain fixed, if you want. An example of why you might want to do this is in the context of helical peptide binding. If you've threaded a helical peptide sequence onto an ideal helix, and now want to diversify the complex, allowing the helix to be predicted now not as an ideal helix, you might do something like:
+
+```
+'contigmap.contigs=[100-100/0 20-20]' 'contigmap.provide_seq=[100-119]' diffuser.partial_T=10
+```
+In this case, the 20aa chain is the helical peptide. The `contigmap.provide_seq` input is zero-indexed, and you can provide a range (so 100-119 is an inclusive range, unmasking the whole sequence of the peptide). Multiple sequence ranges can be provided separated by a comma, e.g. `'contigmap.provide_seq=[172-177,200-205]'`.
+
+Note that the provide_seq option requires using a different model checkpoint, but this is automatically handled by the inference script.
+
+An example of partial diffusion with providing sequence in diffused regions can be found in `./examples/design_partialdiffusion_withseq.sh`. The same example specifying multiple sequence ranges can be found in `./examples/design_partialdiffusion_multipleseq.sh`.
+
+---
+### Binder Design 
+Hopefully, it's now obvious how you might make a binder with diffusion! Indeed, RFdiffusion shows excellent *in silico* and experimental ability to design *de novo* binders. 
+
+<p align="center">
+  <img src="./img/binder.png" alt="alt text" width="950px" align="middle"/>
+</p>
+
+If chain B is your target, then you could do it like this:
+
+```
+./scripts/run_inference.py 'contigmap.contigs=[B1-100/0 100-100]' inference.output_prefix=test_outputs/binder_test inference.num_designs=10
+```
+
+This will generate 100 residue long binders to residues 1-100 of chain B.
+
+However, this probably isn't the best way of making binders. Because diffusion is somewhat computationally-intensive, we need to try and make it as fast as possible. Providing the whole of your target, uncropped, is going to make diffusion very slow if your target is big (and most targets-of-interest, such as cell-surface receptors tend to be *very* big). One tried-and-true method to speed up binder design is to crop the target protein around the desired interface location. BUT! This creates a problem: if you crop your target and potentially expose hydrophobic core residues which were buried before the crop, how can you guarantee the binder will go to the intended interface site on the surface of the target, and not target the tantalizing hydrophobic patch you have just artificially created?
+
+We solve this issue by providing the model with what we call "hotspot residues". The complex models we refer to earlier in this README file have all been trained with hotspot residues, in this training regime, during each example, the model is told (some of) the residues on the target protein which contact the target (i.e., resides that are part of the interface). The model readily learns that it should be making an interface which involved these hotspot residues. At inference time then, we can provide our own hotspot residues to define a region which the binder must contact. These are specified like this: `'ppi.hotspot_res=[A30,A33,A34]'`, where `A` is the chain ID in the input pdb file of the hotspot residue and the number is the residue index in the input pdb file of the hotspot residue.
+
+Finally, it has been observed that the default RFdiffusion model often generates mostly helical binders. These have high computational and experimental success rates. However, there may be cases where other kinds of topologies may be desired. For this, we include a "beta" model, which generates a greater diversity of topologies, but has not been extensively experimentally validated. Try this at your own risk:
+
+```
+inference.ckpt_override_path=models/Complex_beta_ckpt.pt
+```
+
+An example of binder design with RFdiffusion can be found in `./examples/design_ppi.sh`.
+
+---
+
+## Practical Considerations for Binder Design
+
+RFdiffusion is an extremely powerful binder design tool but it is not magic. In this section we will walk through some common pitfalls in RFdiffusion binder design and offer advice on how to get the most out of this method.
+
+### Selecting a Target Site
+Not every site on a target protein is a good candidate for binder design. For a site to be an attractive candidate for binding it should have >~3 hydrophobic residues for the binder to interact with. Binding to charged polar sites is still quite hard. Binding to sites with glycans close to them is also hard since they often become ordered upon binding and you will take an energetic hit for that. Historically, binder design has also avoided unstructured loops, it is not clear if this is still a requirement as RFdiffusion has been used to bind unstructured peptides which share a lot in common with unstructured loops.
+
+### Truncating your Target Protein
+RFdiffusion scales in runtime as O(N^2) where N is the number of residues in your system. As such, it is a very good idea to truncate large targets so that your computations are not unnecessarily	 expensive. RFdiffusion and all downstream steps (including AF2) are designed to allow for a truncated target. Truncating a target is an art. For some targets, such as multidomain extracellular membranes, a natural truncation point is where two domains are joined by a flexible linker. For other proteins, such as virus spike proteins, this truncation point is less obvious. Generally you want to preserve secondary structure and introduce as few chain breaks as possible. You should also try to leave ~10A of target protein on each side of your intended target site. We recommend using PyMol to truncate your target protein.
+
+### Picking Hotspots
+Hotspots are a feature that we integrated into the model to allow for the control of the site on the target which the binder will interact with. In the paper we define a hotspot as a residue on the target protein which is within 10A Cbeta distance of the binder. Of all of the hotspots which are identified on the target 0-20% of these hotspots are actually provided to the model and the rest are masked. This is important for understanding how you should pick hotspots at inference time.; the model is expecting to have to make more contacts than you specify. We normally recommend between 3-6 hotspots, you should run a few pilot runs before generating thousands of designs to make sure the number of hotspots you are providing will give results you like.
+
+If you have run the previous PatchDock RifDock binder design pipeline, for the RFdiffusion paper we chose our hotspots to be the PatchDock residues of the target.
+
+### Binder Design Scale
+In the paper, we generated ~10,000 RFdiffusion binder backbones for each target. From this set of backbones we then generated two sequences per backbone using ProteinMPNN-FastRelax (described below). We screened these ~20,000 designs using AF2 with initial guess and target templating (also described below).
+
+Given the high success rates we observed in the paper, for some targets it may be sufficient to only generate ~1,000 RFdiffusion backbones in a campaign. What you want is to get enough designs that pass pAE_interaction < 10 (described more in Binder Design Filtering section) such that you are able to fill a DNA order with these successful designs. We have found that designs that do not pass pAE_interaction < 10 are not worth ordering since they will likely not work experimentally.
+
+### Sequence Design for Binders
+You may have noticed that the binders designed by RFdiffusion come out with a poly-Glycine sequence. This is not a bug. RFdiffusion is a backbone-generation model and does not generate sequence for the designed region, therefore, another method must be used to assign a sequence to the binders. In the paper we use the ProteinMPNN-FastRelax protocol to do sequence design. We recommend that you do this as well.  The code for this protocol can be found in [this GitHub repo](https://github.com/nrbennet/dl_binder_design). While we did not find the FastRelax part of the protocol to yield the large in silico success rate improvements that it yielded with the RifDock-generated docks, it is still a good way to increase your number of shots-on-goal for each (computationally expensive) RFdiffusion backbone. If you would prefer to simply run ProteinMPNN on your binders without the FastRelax step, that will work fine but will be more computationally expensive.
+
+### Binder Design Filtering
+One of the most important parts of the binder design pipeline is a filtering step to evaluate if your binders are actually predicted to work. In the paper we filtered using AF2 with an initial guess and target templating, scripts for this protocol are available [here](https://github.com/nrbennet/dl_binder_design). We have found that filtering at pae_interaction < 10 is a good predictor of a binder working experimentally.
+
+---
+
+### Fold Conditioning 
+Something that works really well is conditioning binder design (or monomer generation) on particular topologies. This is achieved by providing (partial) secondary structure and block adjacency information (to a model that has been trained to condition on this). 
+<p align="center">
+  <img src="./img/fold_cond.png" alt="alt text" width="950px" align="middle"/>
+</p>
+We are still working out the best way to actually generate this input at inference time, but for now, we have settled upon generating inputs directly from pdb structures. This permits 'low resolution' specification of output topology (i.e., I want a TIM barrel but I don't care precisely where resides are). In `helper_scripts/`, there's a script called `make_secstruc_adj.py`, which can be used as follows:
+
+e.g. 1:
+```
+./make_secstruc_adj.py --input_pdb ./2KL8.pdb --out_dir /my/dir/for/adj_secstruct
+```
+or e.g. 2:
+```
+./make_secstruc_adj.py --pdb_dir ./pdbs/ --out_dir /my/dir/for/adj_secstruct
+```
+
+This will process either a single pdb, or a folder of pdbs, and output a secondary structure and adjacency pytorch file, ready to go into the model. For now (although this might not be necessary), you should also generate these files for the target protein (if you're doing PPI), and provide this to the model. You can then use these at inference as follows:
+
+```
+./scripts/run_inference.py inference.output_prefix=./scaffold_conditioned_test/test scaffoldguided.scaffoldguided=True scaffoldguided.target_pdb=False scaffoldguided.scaffold_dir=./examples/ppi_scaffolds_subset
+```
+
+A few extra things:
+1) As mentioned above, for PPI, you will want to provide a target protein, along with its secondary structure and block adjacency. This can be done by adding:
+
+```
+scaffoldguided.target_pdb=True scaffoldguided.target_path=input_pdbs/insulin_target.pdb inference.output_prefix=insulin_binder/jordi_ss_insulin_noise0_job0 'ppi.hotspot_res=[A59,A83,A91]' scaffoldguided.target_ss=target_folds/insulin_target_ss.pt scaffoldguided.target_adj=target_folds/insulin_target_adj.pt
+```
+
+To generate these block adjacency and secondary structure inputs, you can use the helper script.
+
+This will now generate 3-helix bundles to the insulin target. 
+
+For ppi, it's probably also worth adding this flag:
+
+```
+scaffoldguided.mask_loops=False
+```
+
+This is quite important to understand. During training, we mask some of the secondary structure and block adjacency. This is convenient, because it allows us to, at inference, easily add extra residues without having to specify precise secondary structure for every residue. E.g. if you want to make a long 3 helix bundle, you could mask the loops, and add e.g. 20 more 'mask' tokens to that loop. The model will then (presumbly) choose to make e.g. 15 of these residues into helices (to extend the 3HB), and then make a 5aa loop. But, you didn't have to specify that, which is nice. The way this would be done would be like this:
+
+```
+scaffoldguided.mask_loops=True scaffoldguided.sampled_insertion=15 scaffoldguided.sampled_N=5 scaffoldguided.sampled_C=5
+```
+
+This will, at each run of inference, sample up to 15 residues to insert into loops in your 3HB input, and up to 5 additional residues at N and C terminus.
+This strategy is very useful if you don't have a large set of pdbs to make block adjacencies for. For example, we showed that we could generate loads of lengthened TIM barrels from a single starting pdb with this strategy. However, for PPI, if you're using the provided scaffold sets, it shouldn't be necessary (because there are so many scaffolds to start from, generating extra diversity isn't especially necessary).
+
+Finally, if you have a big directory of block adjacency/secondary structure files, but don't want to use all of them, you can make a `.txt` file of the ones you want to use, and pass:
+
+```
+scaffoldguided.scaffold_list=path/to/list
+```
+
+For PPI, we've consistently seen that reducing the noise added at inference improves designs. This comes at the expense of diversity, but, given that the scaffold sets are huge, this probably doesn't matter too much. We therefore recommend lowering the noise. 0.5 is probably a good compromise:
+
+```
+denoiser.noise_scale_ca=0.5 denoiser.noise_scale_frame=0.5
+```
+This just scales the amount of noise we add to the translations (`noise_scale_ca`) and rotations (`noise_scale_frame`) by, in this case, 0.5.
+
+An additional example of PPI with fold conditioning is available here: `./examples/design_ppi_scaffolded.sh`
+
+---
+
+### Generation of Symmetric Oligomers 
+We're going to switch gears from discussing PPI and look at another task at which RFdiffusion performs well on: symmetric oligomer design. This is done by symmetrising the noise we sample at t=T, and symmetrising the input at every timestep. We have currently implemented the following for use (with the others coming soon!):
+- Cyclic symmetry
+- Dihedral symmetry
+- Tetrahedral symmetry
+
+<p align="center">
+  <img src="./img/olig2.png" alt="alt text" width="1000px" align="middle"/>
+</p>
+
+Here's an example:
+```
+./scripts/run_inference.py --config-name symmetry  inference.symmetry=tetrahedral 'contigmap.contigs=[360]' inference.output_prefix=test_sample/tetrahedral inference.num_designs=1
+```
+
+Here, we've specified a different `config` file (with `--config-name symmetry`). Because symmetric diffusion is quite different from the diffusion described above, we packaged a whole load of symmetry-related configs into a new file (see `configs/inference/symmetry.yml`). Using this config file now puts diffusion in `symmetry-mode`.
+
+The symmetry type is then specified with `inference.symmetry=`. Here, we're specifying tetrahedral symmetry, but you could also choose cyclic (e.g. `c4`) or dihedral (e.g. `d2`).
+
+The configmap.contigs length refers to the *total* length of your oligomer. Therefore, it *must* be divisible by *n* chains.
+
+More examples of designing oligomers can be found here: `./examples/design_cyclic_oligos.sh`, `./examples/design_dihedral_oligos.sh`, `./examples/design_tetrahedral_oligos.sh`.
+
+---
+
+### Using Auxiliary Potentials 
+Performing diffusion with symmetrized noise may give you the idea that we could use other external interventions during the denoising process to guide diffusion. One such intervention that we have implemented is auxiliary potentials. Auxiliary potentials can be very useful for guiding the inference process. E.g. whereas in RFjoint inpainting, we have little/no control over the final shape of an output, in diffusion we can readily force the network to make, for example, a well-packed protein.
+This is achieved in the updates we make at each step.
+
+Let's go a little deeper into how the diffusion process works:
+At timestep T (the first step of the reverse-diffusion inference process), we sample noise from a known *prior* distribution. The model then makes a prediction of what the final structure should be, and we use these two states (noise at time T, prediction of the structure at time 0) to back-calculate where t=T-1 would have been. We therefore have a vector pointing from each coordinate at time T, to their corresponding, back-calculated position at time T-1.
+But, we want to be able to bias this update, to *push* the trajectory towards some desired state. This can be done by biasing that vector with another vector, which points towards a position where that residue would *reduce* the 'loss' as defined by your potential. E.g. if we want to use the `monomer_ROG` potential, which seeks to minimise the radius of gyration of the final protein, if the models prediction of t=0 is very elongated, each of those distant residues will have a larger gradient when we differentiate the `monomer_ROG` potential w.r.t. their positions. These gradients, along with the corresponding scale, can be combined into a vector, which is then combined with the original update vector to make a "biased update" at that timestep.
+
+The exact parameters used when applying these potentials matter. If you weight them too strongly, you're not going to end up with a good protein. Too weak, and they'll have little effect. We've explored these potentials in a few different scenarios, and have set sensible defaults, if you want to use them. But, if you feel like they're too weak/strong, or you just fancy exploring, do play with the parameters (in the `potentials` part of the config file).
+
+Potentials are specified as a list of strings with each string corresponding to a potential. The argument for potentials is `potentials.guiding_potentials`. Within the string per-potential arguments may be specified in the following syntax: `arg_name1:arg_value1,arg_name2:arg_value2,...,arg_nameN:arg_valueN`. The only argument that is required for each potential is the name of the potential that you wish to apply, the name of this argument is `type` as-in the type of potential you wish to use. Some potentials such as `olig_contacts` and `substrate_contacts` take global options such as `potentials.substrate`, see `config/inference/base.yml` for all the global arguments associated with potentials. Additionally, it is useful to have the effect of the potential "decay" throughout the trajectory, such that in the beginning the effect of the potential is 1x strength, and by the end is much weaker. These decays (`constant`,`linear`,`quadratic`,`cubic`) can be set with the `potentials.guide_decay` argument. 
+
+Here's an example of how to specify a potential:
+
+```
+potentials.guiding_potentials=[\"type:olig_contacts,weight_intra:1,weight_inter:0.1\"] potentials.olig_intra_all=True potentials.olig_inter_all=True potentials.guide_scale=2 potentials.guide_decay='quadratic'
+```
+
+We are still fully characterising how/when to use potentials, and we strongly recommend exploring different parameters yourself, as they are clearly somewhat case-dependent. So far, it is clear that they can be helpful for motif scaffolding and symmetric oligomer generation. However, they seem to interact weirdly with hotspot residues in PPI. We think we know why this is, and will work in the coming months to write better potentials for PPI. And please note, it is often good practice to start with *no potentials* as a baseline, then slowly increase their strength. For the oligomer contacts potentials, start with the ones provided in the examples, and note that the `intra` chain potential often should be higher than the `inter` chain potential.
+
+We have already implemented several potentials but it is relatively straightforward to add more, if you want to push your designs towards some specified goal. The *only* condition is that, whatever potential you write, it is differentiable. Take a look at `potentials.potentials.py` for examples of the potentials we have implemented so far. 
+
+---
+
+### Symmetric Motif Scaffolding.  
+We can also combine symmetric diffusion with motif scaffolding to scaffold motifs symmetrically.
+Currently, we have one way for performing symmetric motif scaffolding. That is by specifying the position of the motif specified w.r.t. the symmetry axes.  
+
+<p align="center">
+  <img src="./img/sym_motif.png" alt="alt text" width="1000px" align="middle"/>
+</p>
+
+**Special input .pdb and contigs requirements**
+
+For now, we require that a user have a symmetrized version of their motif in their input pdb for symmetric motif scaffolding. There are two main reasons for this. First, the model is trained by centering any motif at the origin, and thus the code also centers motifs at the origin automatically. Therefore, if your motif is not symmetrized, this centering action will result in an asymmetric unit that now has the origin and axes of symmetry running right through it (bad). Secondly, the diffusion code uses a canonical set of symmetry axes (rotation matrices) to propogate the asymmetric unit of a motif. In order to prevent accidentally running diffusion trajectories which are propogating your motif in ways you don't intend, we require that a user symmetrize an input using the RFdiffusion canonical symmetry axes.
+
+**RFdiffusion canonical symmetry axes** 
+| Group      |      Axis     | 
+|:----------:|:-------------:|
+| Cyclic     |  Z |
+| Dihedral (cyclic) |    Z   |
+| Dihedral (flip/reflection) | X |
+
+
+**Example: Inputs for symmetric motif scaffolding with motif position specified w.r.t the symmetry axes.**
+
+This example script `examples/design_nickel.sh` can be used to scaffold the C4 symmetric Nickel binding domains shown in the RFdiffusion paper. It combines many concepts discussed earlier, including symmetric oligomer generation, motif scaffolding, and use of guiding potentials.
+
+Note that the contigs should specify something that is precisely symmetric. Things will break if this is not the case. 
+
+---
+
+### A Note on Model Weights
+
+Because of everything we want diffusion to be able to do, there is not *One Model To Rule Them All*. E.g., if you want to run with secondary structure conditioning, this requires a different model than if you don't. Under the hood, we take care of most of this by default - we parse your input and work out the most appropriate checkpoint.
+This is where the config setup is really useful. The exact model checkpoint used at inference contains in it all of the parameters is was trained with, so we can just populate the config file with those values, such that inference runs as designed.
+If you do want to specify a different checkpoint (if, for example, we train a new model and you want to test it), you just have to make sure it's compatible with what you're doing. E.g. if you try and give secondary structure features to a model that wasn't trained with them, it'll crash.
+
+### Things you might want to play with at inference time
+
+Occasionally, it might good to try an alternative model (for example the active site model, or the beta binder model). These can be specified with `inference.ckpt_override_path`. We do not recommend using these outside of the described use cases, however, as there is not a guarantee they will understand other kinds of inputs.
+
+For a full list of things that are implemented at inference, see the config file (`configs/inference/base.yml` or `configs/inference/symmetry.yml`). Although you can modify everything, this is not recommended unless you know what you're doing.
+Generally, don't change the `model`, `preprocess` or `diffuser` configs. These pertain to how the model was trained, so it's unwise to change how you use the model at inference time.
+However, the parameters below are definitely worth exploring:
+-inference.final_step: This is when we stop the trajectory. We have seen that you can stop early, and the model is already making a good prediction of the final structure. This speeds up inference.
+-denoiser.noise_scale_ca and denoiser.noise_scale_frame: These can be used to reduce the noise used during sampling (as discussed for PPI above). The default is 1 (the same noise added at training), but this can be reduced to e.g. 0.5, or even 0. This actually improves the quality of models coming out of diffusion, but at the expense of diversity. If you're not getting any good outputs, or if your problem is very constrained, you could try reducing the noise. While these parameters can be changed independently (for translations and rotations), we recommend keeping them tied.
+
+### Understanding the output files
+We output several different files.
+1. The `.pdb` file. This is the final prediction out of the model. Note that every designed residue is output as a glycine (as we only designed the backbone), and no sidechains are output. This is because, even though RFdiffusion conditions on sidechains in an input motif, there is no loss applied to these predictions, so they can't strictly be trusted.
+2. The `.trb` file. This contains useful metadata associated with that specific run, including the specific contig used (if length ranges were sampled), as well as the full config used by RFdiffusion. There are also a few other convenient items in this file:
+    - details about mapping (i.e. how residues in the input map to residues in the output)
+        - `con_ref_pdb_idx`/`con_hal_pdb_idx` - These are two arrays including the input pdb indices (in con_ref_pdb_idx), and where they are in the output pdb (in con_hal_pdb_idx). This only contains the chains where inpainting took place (i.e. not any fixed receptor/target chains)
+        - `con_ref_idx0`/`con_hal_idx0` - These are the same as above, but 0 indexed, and without chain information. This is useful for splicing coordinates out (to assess alignment etc).
+        - `inpaint_seq` - This details any residues that were masked during inference.
+3. Trajectory files. By default, we output the full trajectories into the `/traj/` folder. These files can be opened in pymol, as multi-step pdbs. Note that these are ordered in reverse, so the first pdb is technically the last (t=1) prediction made by RFdiffusion during inference. We include both the `pX0` predictions (what the model predicted at each timestep) and the `Xt-1` trajectories (what went into the model at each timestep).
+
+### Docker
+
+We have provided a Dockerfile at `docker/Dockerfile` to help run RFDiffusion on HPC and other container orchestration systems. Follow these steps to build and run the container on your system:
+
+1. Clone this repository with `git clone https://github.com/RosettaCommons/RFdiffusion.git` and then `cd RFdiffusion`
+1. Verify that the Docker daemon is running on your system with `docker info`. You can find Docker installation instructions for Mac, WIndows, and Linux in the [official Docker docs](https://docs.docker.com/get-docker/). You may also consider [Finch](https://github.com/runfinch/finch), the open source client for container development.
+1. Build the container image on your system with `docker build -f docker/Dockerfile -t rfdiffusion .`
+1. Create some folders on your file system with `mkdir $HOME/inputs $HOME/outputs $HOME/models`
+1. Download the RFDiffusion models with `bash scripts/download_models.sh $HOME/models`
+1. Download a test file (or another of your choice) with `wget -P $HOME/inputs https://files.rcsb.org/view/5TPN.pdb`
+1. Run the container with the following command:
+
+```bash
+docker run -it --rm --gpus all \
+  -v $HOME/models:$HOME/models \
+  -v $HOME/inputs:$HOME/inputs \
+  -v $HOME/outputs:$HOME/outputs \
+  rfdiffusion \
+  inference.output_prefix=$HOME/outputs/motifscaffolding \
+  inference.model_directory_path=$HOME/models \
+  inference.input_pdb=$HOME/inputs/5TPN.pdb \
+  inference.num_designs=3 \
+  'contigmap.contigs=[10-40/A163-181/10-40]'
+```
+
+  This starts the `rfdiffusion` container, mounts the models, inputs, and outputs folders, passes all available GPUs, and then calls the `run_inference.py` script with the parameters specified.
+
+### Conclusion
+
+We are extremely excited to share RFdiffusion with the wider scientific community. We expect to push some updates as and when we make sizeable improvements in the coming months, so do stay tuned. We realize it may take some time to get used to executing RFdiffusion with perfect syntax (sometimes Hydra is hard), so please don't hesitate to create GitHub issues if you need help, we will respond as often as we can. 
+
+Now, let's go make some proteins. Have fun! 
+
+\- Joe, David, Nate, Brian, Jason, and the RFdiffusion team. 
+
+---
+
+RFdiffusion builds directly on the architecture and trained parameters of RoseTTAFold. We therefore thank Frank DiMaio and Minkyung Baek, who developed RoseTTAFold.
+RFdiffusion is released under an open source BSD License (see LICENSE file). It is free for both non-profit and for-profit use. 
+
+

config/inference/base.yaml

@@ -0,0 +1,136 @@
+# Base inference Configuration.
+
+inference:
+  input_pdb: null
+  num_designs: 10
+  design_startnum: 0
+  ckpt_override_path: null
+  symmetry: null
+  recenter: True
+  radius: 10.0
+  model_only_neighbors: False
+  output_prefix: samples/design
+  write_trajectory: True
+  scaffold_guided: False
+  model_runner: SelfConditioning
+  cautious: True
+  align_motif: True
+  symmetric_self_cond: True
+  final_step: 1
+  deterministic: False
+  trb_save_ckpt_path: null
+  schedule_directory_path: null
+  model_directory_path: null
+
+contigmap:
+  contigs: null
+  inpaint_seq: null
+  provide_seq: null
+  length: null
+
+model:
+  n_extra_block: 4
+  n_main_block: 32
+  n_ref_block: 4
+  d_msa: 256
+  d_msa_full: 64
+  d_pair: 128
+  d_templ: 64
+  n_head_msa: 8
+  n_head_pair: 4
+  n_head_templ: 4
+  d_hidden: 32
+  d_hidden_templ: 32
+  p_drop: 0.15
+  SE3_param_full:
+    num_layers: 1
+    num_channels: 32
+    num_degrees: 2
+    n_heads: 4
+    div: 4
+    l0_in_features: 8
+    l0_out_features: 8
+    l1_in_features: 3
+    l1_out_features: 2
+    num_edge_features: 32
+  SE3_param_topk:
+    num_layers: 1
+    num_channels: 32
+    num_degrees: 2
+    n_heads: 4
+    div: 4
+    l0_in_features: 64
+    l0_out_features: 64
+    l1_in_features: 3
+    l1_out_features: 2
+    num_edge_features: 64
+  freeze_track_motif: False
+  use_motif_timestep: False
+
+diffuser:
+  T: 50
+  b_0: 1e-2
+  b_T: 7e-2
+  schedule_type: linear
+  so3_type: igso3
+  crd_scale: 0.25
+  partial_T: null    
+  so3_schedule_type: linear
+  min_b: 1.5
+  max_b: 2.5
+  min_sigma: 0.02
+  max_sigma: 1.5
+
+denoiser:
+  noise_scale_ca: 1
+  final_noise_scale_ca: 1
+  ca_noise_schedule_type: constant
+  noise_scale_frame: 1
+  final_noise_scale_frame: 1
+  frame_noise_schedule_type: constant
+
+ppi:
+  hotspot_res: null
+
+potentials:
+  guiding_potentials: null 
+  guide_scale: 10
+  guide_decay: constant
+  olig_inter_all : null
+  olig_intra_all : null
+  olig_custom_contact : null
+  substrate: null
+
+contig_settings:
+  ref_idx: null
+  hal_idx: null
+  idx_rf: null
+  inpaint_seq_tensor: null
+
+preprocess:
+  sidechain_input: False
+  motif_sidechain_input: True
+  d_t1d: 22
+  d_t2d: 44
+  prob_self_cond: 0.0
+  str_self_cond: False
+  predict_previous: False
+  
+logging:
+  inputs: False
+
+scaffoldguided:
+  scaffoldguided: False
+  target_pdb: False
+  target_path: null
+  scaffold_list: null
+  scaffold_dir: null
+  sampled_insertion: 0
+  sampled_N: 0
+  sampled_C: 0
+  ss_mask: 0
+  systematic: False
+  target_ss: null
+  target_adj: null
+  mask_loops: True
+  contig_crop: null

config/inference/symmetry.yaml

@@ -0,0 +1,26 @@
+# Config for sampling symmetric assemblies.
+
+defaults:
+  - base
+
+inference:
+  # Symmetry to sample
+  # Available symmetries:
+  # - Cyclic symmetry (C_n) # call as c5
+  # - Dihedral symmetry (D_n) # call as d5
+  # - Tetrahedral symmetry # call as tetrahedral
+  # - Octahedral symmetry # call as octahedral
+  # - Icosahedral symmetry # call as icosahedral
+  symmetry: c2
+
+  # Set to true for computational efficiency 
+  # to avoid memory overhead of modeling all subunits.
+  model_only_neighbors: False
+
+  # Output directory of samples.
+  output_prefix: samples/c2
+
+contigmap:
+  # Specify a single integer value to sample unconditionally.
+  # Must be evenly divisible by the number of chains in the symmetry.
+  contigs: ['100']

docker/Dockerfile

@@ -0,0 +1,50 @@
+# Usage: 
+# git clone https://github.com/RosettaCommons/RFdiffusion.git
+# cd RFdiffusion
+# docker build -f docker/Dockerfile -t rfdiffusion .
+# mkdir $HOME/inputs $HOME/outputs $HOME/models
+# bash scripts/download_models.sh $HOME/models
+# wget -P $HOME/inputs https://files.rcsb.org/view/5TPN.pdb
+
+# docker run -it --rm --gpus all \
+#   -v $HOME/models:$HOME/models \
+#   -v $HOME/inputs:$HOME/inputs \
+#   -v $HOME/outputs:$HOME/outputs \
+#   rfdiffusion \
+#   inference.output_prefix=$HOME/outputs/motifscaffolding \
+#   inference.model_directory_path=$HOME/models \
+#   inference.input_pdb=$HOME/inputs/5TPN.pdb \
+#   inference.num_designs=3 \
+#   'contigmap.contigs=[10-40/A163-181/10-40]'
+
+FROM nvcr.io/nvidia/cuda:11.6.2-cudnn8-runtime-ubuntu20.04
+
+COPY . /app/RFdiffusion/
+
+RUN apt-get -q update \ 
+  && DEBIAN_FRONTEND=noninteractive apt-get install --no-install-recommends -y \
+  git \
+  python3.9 \
+  python3-pip \
+  && python3.9 -m pip install -q -U --no-cache-dir pip \
+  && rm -rf /var/lib/apt/lists/* \
+  && apt-get autoremove -y \
+  && apt-get clean \
+  && pip install -q --no-cache-dir \
+  dgl==1.0.2+cu116 -f https://data.dgl.ai/wheels/cu116/repo.html \
+  torch==1.12.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116 \
+  e3nn==0.3.3 \
+  wandb==0.12.0 \
+  pynvml==11.0.0 \
+  git+https://github.com/NVIDIA/dllogger#egg=dllogger \
+  decorator==5.1.0 \
+  hydra-core==1.3.2 \
+  pyrsistent==0.19.3 \
+  /app/RFdiffusion/env/SE3Transformer \
+  && pip install --no-cache-dir /app/RFdiffusion --no-deps
+  
+WORKDIR /app/RFdiffusion
+
+ENV DGLBACKEND="pytorch"
+
+ENTRYPOINT ["python3.9", "scripts/run_inference.py"]

env/SE3Transformer/.dockerignore

@@ -0,0 +1,123 @@
+.Trash-0
+.git
+data/
+.DS_Store
+*wandb/
+*.pt
+*.swp
+
+# added by FAFU
+.idea/
+cache/
+downloaded/
+*.lprof
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+**/benchmark
+**/results
+*.pkl
+*.log

env/SE3Transformer/.gitignore

@@ -0,0 +1,121 @@
+data/
+.DS_Store
+*wandb/
+*.pt
+*.swp
+
+# added by FAFU
+.idea/
+cache/
+downloaded/
+*.lprof
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+**/benchmark
+**/results
+*.pkl
+*.log

env/SE3Transformer/Dockerfile

@@ -0,0 +1,58 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+# run docker daemon with --default-runtime=nvidia for GPU detection during build
+# multistage build for DGL with CUDA and FP16
+
+ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:21.07-py3
+
+FROM ${FROM_IMAGE_NAME} AS dgl_builder
+
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update \
+    && apt-get install -y git build-essential python3-dev make cmake \
+    && rm -rf /var/lib/apt/lists/*
+WORKDIR /dgl
+RUN git clone --branch v0.7.0 --recurse-submodules --depth 1 https://github.com/dmlc/dgl.git .
+RUN sed -i 's/"35 50 60 70"/"60 70 80"/g' cmake/modules/CUDA.cmake
+WORKDIR build
+RUN cmake -DUSE_CUDA=ON -DUSE_FP16=ON ..
+RUN make -j8
+
+
+FROM ${FROM_IMAGE_NAME}
+
+RUN rm -rf /workspace/*
+WORKDIR /workspace/se3-transformer
+
+# copy built DGL and install it
+COPY --from=dgl_builder /dgl ./dgl
+RUN cd dgl/python && python setup.py install && cd ../.. && rm -rf dgl
+
+ADD requirements.txt .
+RUN pip install --no-cache-dir --upgrade --pre pip
+RUN pip install --no-cache-dir -r requirements.txt
+ADD . .
+
+ENV DGLBACKEND=pytorch
+ENV OMP_NUM_THREADS=1

env/SE3Transformer/LICENSE

@@ -0,0 +1,7 @@
+Copyright 2021 NVIDIA CORPORATION & AFFILIATES
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

env/SE3Transformer/NOTICE

@@ -0,0 +1,7 @@
+SE(3)-Transformer PyTorch
+
+This repository includes software from https://github.com/FabianFuchsML/se3-transformer-public
+licensed under the MIT License.
+
+This repository includes software from https://github.com/lucidrains/se3-transformer-pytorch
+licensed under the MIT License.

env/SE3Transformer/README.md

@@ -0,0 +1,580 @@
+# SE(3)-Transformers For PyTorch
+
+This repository provides a script and recipe to train the SE(3)-Transformer model to achieve state-of-the-art accuracy. The content of this repository is tested and maintained by NVIDIA.
+
+## Table Of Contents
+- [Model overview](#model-overview)
+    * [Model architecture](#model-architecture)
+    * [Default configuration](#default-configuration)
+    * [Feature support matrix](#feature-support-matrix)
+        * [Features](#features)
+    * [Mixed precision training](#mixed-precision-training)
+        * [Enabling mixed precision](#enabling-mixed-precision)
+          * [Enabling TF32](#enabling-tf32)
+    * [Glossary](#glossary)
+- [Setup](#setup)
+    * [Requirements](#requirements)
+- [Quick Start Guide](#quick-start-guide)
+- [Advanced](#advanced)
+    * [Scripts and sample code](#scripts-and-sample-code)
+    * [Parameters](#parameters)
+    * [Command-line options](#command-line-options)
+    * [Getting the data](#getting-the-data)
+        * [Dataset guidelines](#dataset-guidelines)
+        * [Multi-dataset](#multi-dataset)
+    * [Training process](#training-process)
+    * [Inference process](#inference-process)
+- [Performance](#performance)
+    * [Benchmarking](#benchmarking)
+        * [Training performance benchmark](#training-performance-benchmark)
+        * [Inference performance benchmark](#inference-performance-benchmark)
+    * [Results](#results)
+        * [Training accuracy results](#training-accuracy-results)                         
+            * [Training accuracy: NVIDIA DGX A100 (8x A100 80GB)](#training-accuracy-nvidia-dgx-a100-8x-a100-80gb)  
+            * [Training accuracy: NVIDIA DGX-1 (8x V100 16GB)](#training-accuracy-nvidia-dgx-1-8x-v100-16gb)
+            * [Training stability test](#training-stability-test)
+        * [Training performance results](#training-performance-results)
+            * [Training performance: NVIDIA DGX A100 (8x A100 80GB)](#training-performance-nvidia-dgx-a100-8x-a100-80gb) 
+            * [Training performance: NVIDIA DGX-1 (8x V100 16GB)](#training-performance-nvidia-dgx-1-8x-v100-16gb)
+        * [Inference performance results](#inference-performance-results)
+            * [Inference performance: NVIDIA DGX A100 (1x A100 80GB)](#inference-performance-nvidia-dgx-a100-1x-a100-80gb)
+            * [Inference performance: NVIDIA DGX-1 (1x V100 16GB)](#inference-performance-nvidia-dgx-1-1x-v100-16gb)
+- [Release notes](#release-notes)
+    * [Changelog](#changelog)
+    * [Known issues](#known-issues)
+
+
+
+## Model overview
+
+
+The **SE(3)-Transformer** is a Graph Neural Network using a variant of [self-attention](https://arxiv.org/abs/1706.03762v5) for 3D points and graphs processing.
+This model is [equivariant](https://en.wikipedia.org/wiki/Equivariant_map) under [continuous 3D roto-translations](https://en.wikipedia.org/wiki/Euclidean_group), meaning that when the inputs (graphs or sets of points) rotate in 3D space (or more generally experience a [proper rigid transformation](https://en.wikipedia.org/wiki/Rigid_transformation)), the model outputs either stay invariant or transform with the input.
+A mathematical guarantee of equivariance is important to ensure stable and predictable performance in the presence of nuisance transformations of the data input and when the problem has some inherent symmetries we want to exploit.
+
+
+The model is based on the following publications:
+- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://arxiv.org/abs/2006.10503) (NeurIPS 2020) by Fabian B. Fuchs, Daniel E. Worrall, et al. 
+- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://arxiv.org/abs/1802.08219) by Nathaniel Thomas, Tess Smidt, et al.
+
+A follow-up paper explains how this model can be used iteratively, for example, to predict or refine protein structures:
+
+- [Iterative SE(3)-Transformers](https://arxiv.org/abs/2102.13419) by Fabian B. Fuchs, Daniel E. Worrall, et al. 
+
+Just like [the official implementation](https://github.com/FabianFuchsML/se3-transformer-public), this implementation uses [PyTorch](https://pytorch.org/) and the [Deep Graph Library (DGL)](https://www.dgl.ai/).
+
+The main differences between this implementation of SE(3)-Transformers and the official one are the following:
+
+- Training and inference support for multiple GPUs
+- Training and inference support for [Mixed Precision](https://arxiv.org/abs/1710.03740)
+- The [QM9 dataset from DGL](https://docs.dgl.ai/en/latest/api/python/dgl.data.html#qm9edge-dataset) is used and automatically downloaded
+- Significantly increased throughput
+- Significantly reduced memory consumption
+- The use of layer normalization in the fully connected radial profile layers is an option (`--use_layer_norm`), off by default
+- The use of equivariant normalization between attention layers is an option (`--norm`), off by default
+- The [spherical harmonics](https://en.wikipedia.org/wiki/Spherical_harmonic) and [Clebsch–Gordan coefficients](https://en.wikipedia.org/wiki/Clebsch%E2%80%93Gordan_coefficients), used to compute bases matrices, are computed with the [e3nn library](https://e3nn.org/)
+
+
+
+This model enables you to predict quantum chemical properties of small organic molecules in the [QM9 dataset](https://www.nature.com/articles/sdata201422).
+In this case, the exploited symmetry is that these properties do not depend on the orientation or position of the molecules in space.
+
+
+This model is trained with mixed precision using Tensor Cores on NVIDIA Volta, NVIDIA Turing, and the NVIDIA Ampere GPU architectures. Therefore, researchers can get results up to 1.5x faster than training without Tensor Cores while experiencing the benefits of mixed precision training. This model is tested against each NGC monthly container release to ensure consistent accuracy and performance over time.
+
+### Model architecture
+
+The model consists of stacked layers of equivariant graph self-attention and equivariant normalization.
+Lastly, a Tensor Field Network convolution is applied to obtain invariant features. Graph pooling (mean or max over the nodes) is applied to these features, and the result is fed to a final MLP to get scalar predictions.
+
+In this setup, the model is a graph-to-scalar network. The pooling can be removed to obtain a graph-to-graph network, and the final TFN can be modified to output features of any type (invariant scalars, 3D vectors, ...).
+
+
+![Model high-level architecture](./images/se3-transformer.png)
+
+
+### Default configuration
+
+
+SE(3)-Transformers introduce a self-attention layer for graphs that is equivariant to 3D roto-translations. It achieves this by leveraging Tensor Field Networks to build attention weights that are invariant and attention values that are equivariant.
+Combining the equivariant values with the invariant weights gives rise to an equivariant output. This output is normalized while preserving equivariance thanks to equivariant normalization layers operating on feature norms.
+
+
+The following features were implemented in this model:
+
+- Support for edge features of any degree (1D, 3D, 5D, ...), whereas the official implementation only supports scalar invariant edge features (degree 0). Edge features with a degree greater than one are
+concatenated to node features of the same degree. This is required in order to reproduce published results on point cloud processing.
+- Data-parallel multi-GPU training (DDP)
+- Mixed precision training (autocast, gradient scaling)
+- Gradient accumulation
+- Model checkpointing
+
+
+The following performance optimizations were implemented in this model:
+
+
+**General optimizations**
+
+- The option is provided to precompute bases at the beginning of the training instead of computing them at the beginning of each forward pass (`--precompute_bases`)
+- The bases computation is just-in-time (JIT) compiled with `torch.jit.script`
+- The Clebsch-Gordon coefficients are cached in RAM
+
+
+**Tensor Field Network optimizations**
+
+- The last layer of each radial profile network does not add any bias in order to avoid large broadcasting operations
+- The layout (order of dimensions) of the bases tensors is optimized to avoid copies to contiguous memory in the downstream TFN layers
+- When Tensor Cores are available, and the output feature dimension of computed bases is odd, then it is padded with zeros to make more effective use of Tensor Cores (AMP and TF32 precisions)
+- Multiple levels of fusion for TFN convolutions (and radial profiles) are provided and automatically used when conditions are met
+- A low-memory mode is provided that will trade throughput for less memory use (`--low_memory`)
+
+**Self-attention optimizations**
+
+- Attention keys and values are computed by a single partial TFN graph convolution in each attention layer instead of two
+- Graph operations for different output degrees may be fused together if conditions are met
+
+
+**Normalization optimizations**
+
+- The equivariant normalization layer is optimized from multiple layer normalizations to a group normalization on fused norms when certain conditions are met
+    
+
+
+Competitive training results and analysis are provided for the following hyperparameters (identical to the ones in the original publication):
+- Number of layers: 7
+- Number of degrees: 4
+- Number of channels: 32
+- Number of attention heads: 8
+- Channels division: 2
+- Use of equivariant normalization: true
+- Use of layer normalization: true
+- Pooling: max
+
+
+### Feature support matrix
+
+This model supports the following features:: 
+
+| Feature               | SE(3)-Transformer                
+|-----------------------|--------------------------
+|Automatic mixed precision (AMP)   |         Yes 
+|Distributed data parallel (DDP)   |         Yes 
+         
+#### Features
+
+
+**Distributed data parallel (DDP)**
+
+[DistributedDataParallel (DDP)](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) implements data parallelism at the module level that can run across multiple GPUs or machines.
+
+**Automatic Mixed Precision (AMP)**
+
+This implementation uses the native PyTorch AMP implementation of mixed precision training. It allows us to use FP16 training with FP32 master weights by modifying just a few lines of code. A detailed explanation of mixed precision can be found in the next section.
+
+### Mixed precision training
+
+Mixed precision is the combined use of different numerical precisions in a computational method. [Mixed precision](https://arxiv.org/abs/1710.03740) training offers significant computational speedup by performing operations in half-precision format while storing minimal information in single-precision to retain as much information as possible in critical parts of the network. Since the introduction of [Tensor Cores](https://developer.nvidia.com/tensor-cores) in NVIDIA Volta, and following with both the NVIDIA Turing and NVIDIA Ampere Architectures, significant training speedups are experienced by switching to mixed precision -- up to 3x overall speedup on the most arithmetically intense model architectures. Using [mixed precision training](https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html) previously required two steps:
+1.  Porting the model to use the FP16 data type where appropriate.    
+2.  Adding loss scaling to preserve small gradient values.
+
+AMP enables mixed precision training on NVIDIA Volta, NVIDIA Turing, and NVIDIA Ampere GPU architectures automatically. The PyTorch framework code makes all necessary model changes internally.
+
+For information about:
+-   How to train using mixed precision, refer to the [Mixed Precision Training](https://arxiv.org/abs/1710.03740) paper and [Training With Mixed Precision](https://docs.nvidia.com/deeplearning/performance/mixed-precision-training/index.html) documentation.
+-   Techniques used for mixed precision training, refer to the [Mixed-Precision Training of Deep Neural Networks](https://devblogs.nvidia.com/mixed-precision-training-deep-neural-networks/) blog.
+-   APEX tools for mixed precision training, refer to the [NVIDIA Apex: Tools for Easy Mixed-Precision Training in PyTorch](https://devblogs.nvidia.com/apex-pytorch-easy-mixed-precision-training/).
+
+#### Enabling mixed precision
+
+Mixed precision is enabled in PyTorch by using the native [Automatic Mixed Precision package](https://pytorch.org/docs/stable/amp.html), which casts variables to half-precision upon retrieval while storing variables in single-precision format. Furthermore, to preserve small gradient magnitudes in backpropagation, a [loss scaling](https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/index.html#lossscaling) step must be included when applying gradients. In PyTorch, loss scaling can be applied automatically using a `GradScaler`.
+Automatic Mixed Precision makes all the adjustments internally in PyTorch, providing two benefits over manual operations. First, programmers need not modify network model code, reducing development and maintenance effort. Second, using AMP maintains forward and backward compatibility with all the APIs for defining and running PyTorch models.
+
+To enable mixed precision, you can simply use the `--amp` flag when running the training or inference scripts.
+
+#### Enabling TF32
+
+TensorFloat-32 (TF32) is the new math mode in [NVIDIA A100](https://www.nvidia.com/en-us/data-center/a100/) GPUs for handling the matrix math, also called tensor operations. TF32 running on Tensor Cores in A100 GPUs can provide up to 10x speedups compared to single-precision floating-point math (FP32) on NVIDIA Volta GPUs. 
+
+TF32 Tensor Cores can speed up networks using FP32, typically with no loss of accuracy. It is more robust than FP16 for models that require a high dynamic range for weights or activations.
+
+For more information, refer to the [TensorFloat-32 in the A100 GPU Accelerates AI Training, HPC up to 20x](https://blogs.nvidia.com/blog/2020/05/14/tensorfloat-32-precision-format/) blog post.
+
+TF32 is supported in the NVIDIA Ampere GPU architecture and is enabled by default.
+
+
+
+### Glossary
+
+**Degree (type)**
+
+In the model, every feature (input, output and hidden) transforms in an equivariant way in relation to the input graph. When we define a feature, we need to choose, in addition to the number of channels, which transformation rule it obeys.
+
+The degree or type of a feature is a positive integer that describes how this feature transforms when the input rotates in 3D.
+
+This is related to [irreducible representations](https://en.wikipedia.org/wiki/Irreducible_representation) of different rotation orders.
+
+The degree of a feature determines its dimensionality. A type-d feature has a dimensionality of 2d+1.
+
+Some common examples include:
+- Degree 0: 1D scalars invariant to rotation
+- Degree 1: 3D vectors that rotate according to 3D rotation matrices
+- Degree 2: 5D vectors that rotate according to 5D [Wigner-D matrices](https://en.wikipedia.org/wiki/Wigner_D-matrix). These can represent symmetric traceless 3x3 matrices.
+
+**Fiber**
+
+A fiber can be viewed as a representation of a set of features of different types or degrees (positive integers), where each feature type transforms according to its rule.
+
+In this repository, a fiber can be seen as a dictionary with degrees as keys and numbers of channels as values.
+
+**Multiplicity**
+
+The multiplicity of a feature of a given type is the number of channels of this feature.
+
+**Tensor Field Network**
+
+A [Tensor Field Network](https://arxiv.org/abs/1802.08219) is a kind of equivariant graph convolution that can combine features of different degrees and produce new ones while preserving equivariance thanks to [tensor products](https://en.wikipedia.org/wiki/Tensor_product).
+
+**Equivariance**
+
+[Equivariance](https://en.wikipedia.org/wiki/Equivariant_map) is a property of a function of model stating that applying a symmetry transformation to the input and then computing the function produces the same result as computing the function and then applying the transformation to the output.
+
+In the case of SE(3)-Transformer, the symmetry group is the group of continuous roto-translations (SE(3)).
+
+## Setup
+
+The following section lists the requirements that you need to meet in order to start training the SE(3)-Transformer model.
+
+### Requirements
+
+This repository contains a Dockerfile which extends the PyTorch 21.07 NGC container and encapsulates some dependencies. Aside from these dependencies, ensure you have the following components:
+- [NVIDIA Docker](https://github.com/NVIDIA/nvidia-docker)
+- PyTorch 21.07+ NGC container
+- Supported GPUs:
+    - [NVIDIA Volta architecture](https://www.nvidia.com/en-us/data-center/volta-gpu-architecture/)
+    - [NVIDIA Turing architecture](https://www.nvidia.com/en-us/design-visualization/technologies/turing-architecture/)
+    - [NVIDIA Ampere architecture](https://www.nvidia.com/en-us/data-center/nvidia-ampere-gpu-architecture/)
+
+For more information about how to get started with NGC containers, refer to the following sections from the NVIDIA GPU Cloud Documentation and the Deep Learning Documentation:
+- [Getting Started Using NVIDIA GPU Cloud](https://docs.nvidia.com/ngc/ngc-getting-started-guide/index.html)
+- [Accessing And Pulling From The NGC Container Registry](https://docs.nvidia.com/deeplearning/frameworks/user-guide/index.html#accessing_registry)
+- [Running PyTorch](https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/running.html#running)
+  
+For those unable to use the PyTorch NGC container to set up the required environment or create your own container, refer to the versioned [NVIDIA Container Support Matrix](https://docs.nvidia.com/deeplearning/frameworks/support-matrix/index.html).
+
+## Quick Start Guide
+
+To train your model using mixed or TF32 precision with Tensor Cores or FP32, perform the following steps using the default parameters of the SE(3)-Transformer model on the QM9 dataset. For the specifics concerning training and inference, refer to the [Advanced](#advanced) section.
+
+1. Clone the repository.
+    ```
+    git clone https://github.com/NVIDIA/DeepLearningExamples
+    cd DeepLearningExamples/PyTorch/DrugDiscovery/SE3Transformer
+    ```
+   
+2.  Build the `se3-transformer` PyTorch NGC container.
+    ```
+    docker build -t se3-transformer .
+    ```
+
+3.  Start an interactive session in the NGC container to run training/inference.
+    ```
+    mkdir -p results
+    docker run -it --runtime=nvidia --shm-size=8g --ulimit memlock=-1 --ulimit stack=67108864 --rm -v ${PWD}/results:/results se3-transformer:latest
+    ```
+
+4. Start training.
+   ```
+   bash scripts/train.sh
+   ```
+
+5. Start inference/predictions.
+   ```
+   bash scripts/predict.sh
+   ```
+
+
+Now that you have your model trained and evaluated, you can choose to compare your training results with our [Training accuracy results](#training-accuracy-results). You can also choose to benchmark your performance to [Training performance benchmark](#training-performance-results) or [Inference performance benchmark](#inference-performance-results). Following the steps in these sections will ensure that you achieve the same accuracy and performance results as stated in the [Results](#results) section.
+
+## Advanced
+
+The following sections provide greater details of the dataset, running training and inference, and the training results.
+
+### Scripts and sample code
+
+In the root directory, the most important files are:
+- `Dockerfile`: container with the basic set of dependencies to run SE(3)-Transformers
+- `requirements.txt`: set of extra requirements to run SE(3)-Transformers
+- `se3_transformer/data_loading/qm9.py`: QM9 data loading and preprocessing, as well as bases precomputation
+- `se3_transformer/model/layers/`: directory containing model architecture layers
+- `se3_transformer/model/transformer.py`: main Transformer module
+- `se3_transformer/model/basis.py`: logic for computing bases matrices
+- `se3_transformer/runtime/training.py`: training script, to be run as a python module
+- `se3_transformer/runtime/inference.py`: inference script, to be run as a python module
+- `se3_transformer/runtime/metrics.py`: MAE metric with support for multi-GPU synchronization
+- `se3_transformer/runtime/loggers.py`: [DLLogger](https://github.com/NVIDIA/dllogger) and [W&B](wandb.ai/) loggers
+
+
+### Parameters
+
+The complete list of the available parameters for the `training.py` script contains:
+
+**General**
+
+- `--epochs`: Number of training epochs (default: `100` for single-GPU)
+- `--batch_size`: Batch size (default: `240`)
+- `--seed`: Set a seed globally (default: `None`)
+- `--num_workers`: Number of dataloading workers (default: `8`)
+- `--amp`: Use Automatic Mixed Precision (default `false`)
+- `--gradient_clip`: Clipping of the gradient norms (default: `None`)
+- `--accumulate_grad_batches`: Gradient accumulation (default: `1`)
+- `--ckpt_interval`: Save a checkpoint every N epochs (default: `-1`)
+- `--eval_interval`: Do an evaluation round every N epochs (default: `1`)
+- `--silent`: Minimize stdout output (default: `false`)
+
+**Paths**
+
+- `--data_dir`: Directory where the data is located or should be downloaded (default: `./data`)
+- `--log_dir`: Directory where the results logs should be saved (default: `/results`)
+- `--save_ckpt_path`: File where the checkpoint should be saved (default: `None`)
+- `--load_ckpt_path`: File of the checkpoint to be loaded (default: `None`)
+
+**Optimizer**
+
+- `--optimizer`: Optimizer to use (default: `adam`)
+- `--learning_rate`: Learning rate to use (default: `0.002` for single-GPU)
+- `--momentum`: Momentum to use (default: `0.9`)
+- `--weight_decay`: Weight decay to use (default: `0.1`)
+
+**QM9 dataset**
+
+- `--task`: Regression task to train on (default: `homo`)
+- `--precompute_bases`: Precompute bases at the beginning of the script during dataset initialization, instead of computing them at the beginning of each forward pass (default: `false`)
+
+**Model architecture**
+
+- `--num_layers`: Number of stacked Transformer layers (default: `7`)
+- `--num_heads`: Number of heads in self-attention (default: `8`)
+- `--channels_div`: Channels division before feeding to attention layer (default: `2`)
+- `--pooling`: Type of graph pooling (default: `max`)
+- `--norm`: Apply a normalization layer after each attention block (default: `false`)
+- `--use_layer_norm`: Apply layer normalization between MLP layers (default: `false`)
+- `--low_memory`: If true, will use fused ops that are slower but use less memory (expect 25 percent less memory). Only has an effect if AMP is enabled on NVIDIA Volta GPUs or if running on Ampere GPUs (default: `false`)
+- `--num_degrees`: Number of degrees to use. Hidden features will have types [0, ..., num_degrees - 1] (default: `4`)
+- `--num_channels`: Number of channels for the hidden features (default: `32`)
+
+
+### Command-line options
+
+To show the full list of available options and their descriptions, use the `-h` or `--help` command-line option, for example: `python -m se3_transformer.runtime.training --help`.
+
+
+### Dataset guidelines
+
+#### Demo dataset
+
+The SE(3)-Transformer was trained on the QM9 dataset.
+
+The QM9 dataset is hosted on DGL servers and downloaded (38MB) automatically when needed. By default, it is stored in the `./data` directory, but this location can be changed with the `--data_dir` argument.
+
+The dataset is saved as a `qm9_edge.npz` file and converted to DGL graphs at runtime.
+
+As input features, we use:
+- Node features (6D):
+    - One-hot-encoded atom type (5D) (atom types: H, C, N, O, F)
+    - Number of protons of each atom (1D)
+- Edge features: one-hot-encoded bond type (4D) (bond types: single, double, triple, aromatic)
+- The relative positions between adjacent nodes (atoms)
+
+#### Custom datasets
+
+To use this network on a new dataset, you can extend the `DataModule` class present in `se3_transformer/data_loading/data_module.py`.
+
+Your custom collate function should return a tuple with:
+
+- A (batched) DGLGraph object
+- A dictionary of node features ({‘{degree}’: tensor})
+- A dictionary of edge features ({‘{degree}’: tensor})
+- (Optional) Precomputed bases as a dictionary
+- Labels as a tensor
+
+You can then modify the `training.py` and `inference.py` scripts to use your new data module.
+
+### Training process
+
+The training script is `se3_transformer/runtime/training.py`, to be run as a module: `python -m se3_transformer.runtime.training`.
+
+**Logs**
+
+By default, the resulting logs are stored in `/results/`. This can be changed with `--log_dir`.
+
+You can connect your existing Weights & Biases account by setting the `WANDB_API_KEY` environment variable.
+
+**Checkpoints**
+
+The argument `--save_ckpt_path` can be set to the path of the file where the checkpoints should be saved.
+`--ckpt_interval` can also be set to the interval (in the number of epochs) between checkpoints.
+
+**Evaluation**
+
+The evaluation metric is the Mean Absolute Error (MAE).
+
+`--eval_interval` can be set to the interval (in the number of epochs) between evaluation rounds. By default, an evaluation round is performed after each epoch.
+
+**Automatic Mixed Precision**
+
+To enable Mixed Precision training, add the `--amp` flag.
+
+**Multi-GPU and multi-node**
+
+The training script supports the PyTorch elastic launcher to run on multiple GPUs or nodes.  Refer to the [official documentation](https://pytorch.org/docs/1.9.0/elastic/run.html).
+
+For example, to train on all available GPUs with AMP:
+
+```
+python -m torch.distributed.run --nnodes=1 --nproc_per_node=gpu --module se3_transformer.runtime.training --amp
+```
+
+
+### Inference process
+
+Inference can be run by using the `se3_transformer.runtime.inference` python module.
+
+The inference script is `se3_transformer/runtime/inference.py`, to be run as a module: `python -m se3_transformer.runtime.inference`.  It requires a pre-trained model checkpoint (to be passed as `--load_ckpt_path`).
+
+
+## Performance
+
+The performance measurements in this document were conducted at the time of publication and may not reflect the performance achieved from NVIDIA’s latest software release. For the most up-to-date performance measurements, go to [NVIDIA Data Center Deep Learning Product Performance](https://developer.nvidia.com/deep-learning-performance-training-inference).
+
+### Benchmarking
+
+The following section shows how to run benchmarks measuring the model performance in training and inference modes.
+
+#### Training performance benchmark
+
+To benchmark the training performance on a specific batch size, run `bash scripts/benchmarck_train.sh {BATCH_SIZE}` for single GPU, and `bash scripts/benchmarck_train_multi_gpu.sh {BATCH_SIZE}` for multi-GPU.
+
+#### Inference performance benchmark
+
+To benchmark the inference performance on a specific batch size, run `bash scripts/benchmarck_inference.sh {BATCH_SIZE}`.
+
+### Results
+
+
+The following sections provide details on how we achieved our performance and accuracy in training and inference.
+
+#### Training accuracy results
+
+##### Training accuracy: NVIDIA DGX A100 (8x A100 80GB)
+
+Our results were obtained by running the `scripts/train.sh` training script in the PyTorch 21.07 NGC container on NVIDIA DGX A100 (8x A100 80GB) GPUs.
+
+| GPUs    | Batch size / GPU    | Absolute error - TF32  | Absolute error - mixed precision  |   Time to train - TF32  |  Time to train - mixed precision | Time to train speedup (mixed precision to TF32) |       
+|:------------------:|:----------------------:|:--------------------:|:------------------------------------:|:---------------------------------:|:----------------------:|:----------------------------------------------:|
+|  1                 |    240                   |           0.03456                            |        0.03460                                |        1h23min      |    1h03min                |    1.32x              |
+|  8                 |    240                   |           0.03417                            |        0.03424                                |        15min          |    12min                |    1.25x              |
+
+
+##### Training accuracy: NVIDIA DGX-1 (8x V100 16GB)
+
+Our results were obtained by running the `scripts/train.sh` training script in the PyTorch 21.07 NGC container on NVIDIA DGX-1 with (8x V100 16GB) GPUs.
+
+| GPUs    | Batch size / GPU    | Absolute error - FP32  | Absolute error - mixed precision  |   Time to train - FP32  |  Time to train - mixed precision | Time to train speedup (mixed precision to FP32)  |      
+|:------------------:|:----------------------:|:--------------------:|:------------------------------------:|:---------------------------------:|:----------------------:|:----------------------------------------------:|
+|  1                 |    240                   |           0.03432                            |        0.03439                                |         2h25min         |    1h33min                |    1.56x              |
+|  8                 |    240                   |           0.03380                            |        0.03495                                |        29min          |    20min                |    1.45x              |
+
+
+#### Training performance results
+
+##### Training performance: NVIDIA DGX A100 (8x A100 80GB)
+
+Our results were obtained by running the `scripts/benchmark_train.sh` and `scripts/benchmark_train_multi_gpu.sh` benchmarking scripts in the PyTorch 21.07 NGC container on NVIDIA DGX A100 with 8x A100 80GB GPUs. Performance numbers (in molecules per millisecond) were averaged over five  entire training epochs after a warmup epoch.
+
+| GPUs             | Batch size / GPU     | Throughput - TF32 [mol/ms]                             | Throughput - mixed precision [mol/ms]      | Throughput speedup (mixed precision - TF32)   | Weak scaling - TF32    | Weak scaling - mixed precision |
+|:------------------:|:----------------------:|:--------------------:|:------------------------------------:|:---------------------------------:|:----------------------:|:----------------------------------------------:|
+|   1              |     240             |   2.21                                       |   2.92                            |   1.32x                         |                      |                                              |
+|   1              |     120              |  1.81                                        |  2.04                             |  1.13x                          |                      |                                              |
+|   8              |     240             |   17.15                                      |     22.95                         |   1.34x                         |   7.76               |    7.86                                     |
+|   8              |     120              |  13.89                                       |    15.62                          |  1.12x                          |       7.67           |    7.66                                       |
+
+
+To achieve these same results, follow the steps in the [Quick Start Guide](#quick-start-guide).
+
+
+##### Training performance: NVIDIA DGX-1 (8x V100 16GB)
+
+Our results were obtained by running the `scripts/benchmark_train.sh` and `scripts/benchmark_train_multi_gpu.sh` benchmarking scripts in the PyTorch 21.07 NGC container on NVIDIA DGX-1 with 8x V100 16GB GPUs. Performance numbers (in molecules per millisecond) were averaged over five  entire training epochs after a warmup epoch.
+
+| GPUs             | Batch size / GPU     | Throughput - FP32 [mol/ms] | Throughput - mixed precision  [mol/ms]     | Throughput speedup (FP32 - mixed precision)   | Weak scaling - FP32    | Weak scaling - mixed precision |
+|:------------------:|:----------------------:|:--------------------:|:------------------------------------:|:---------------------------------:|:----------------------:|:----------------------------------------------:|
+|   1              |     240              |    1.25          |    1.88                           |  1.50x                          |                      |                                              |
+|   1              |     120              |    1.03           |   1.41                            |  1.37x                          |                      |                                              |
+|   8              |     240              |    9.33           |   14.02                           |  1.50x                          |      7.46            |      7.46                                    |
+|   8              |     120              |    7.39           |   9.41                           |   1.27x                         |        7.17          |        6.67                                  |
+
+
+To achieve these same results, follow the steps in the [Quick Start Guide](#quick-start-guide).
+
+
+#### Inference performance results
+
+
+##### Inference performance: NVIDIA DGX A100 (1x A100 80GB)
+
+Our results were obtained by running the `scripts/benchmark_inference.sh` inferencing benchmarking script in the PyTorch 21.07 NGC container on NVIDIA DGX A100 with 1x A100 80GB GPU.
+
+FP16
+
+| Batch size | Throughput Avg [mol/ms] | Latency Avg [ms] | Latency 90% [ms] |Latency 95% [ms] |Latency 99% [ms] |
+|:------------:|:------:|:-----:|:-----:|:-----:|:-----:|
+| 1600 | 11.60 | 140.94 | 138.29 | 140.12 | 386.40 |
+| 800 | 10.74 | 75.69 | 75.74 | 76.50 | 79.77 |
+| 400 | 8.86 | 45.57 | 46.11 | 46.60 | 49.97 |
+
+TF32
+
+| Batch size | Throughput Avg [mol/ms] | Latency Avg [ms] | Latency 90% [ms] |Latency 95% [ms] |Latency 99% [ms] |
+|:------------:|:------:|:-----:|:-----:|:-----:|:-----:|
+| 1600 | 8.58 | 189.20 | 186.39 | 187.71 | 420.28 |
+| 800 | 8.28 | 97.56 | 97.20 | 97.73 | 101.13 |
+| 400 | 7.55 | 53.38 | 53.72 | 54.48 | 56.62 |
+
+To achieve these same results, follow the steps in the [Quick Start Guide](#quick-start-guide).
+
+
+
+##### Inference performance: NVIDIA DGX-1 (1x V100 16GB)
+
+Our results were obtained by running the `scripts/benchmark_inference.sh` inferencing benchmarking script in the PyTorch 21.07 NGC container on NVIDIA DGX-1 with 1x V100 16GB GPU.
+
+FP16
+
+| Batch size | Throughput Avg [mol/ms] | Latency Avg [ms] | Latency 90% [ms] |Latency 95% [ms] |Latency 99% [ms] |
+|:------------:|:------:|:-----:|:-----:|:-----:|:-----:|
+| 1600 | 6.42 | 254.54 | 247.97 | 249.29 | 721.15 |
+| 800 | 6.13 | 132.07 | 131.90 | 132.70 | 140.15 |
+| 400 | 5.37 | 75.12 | 76.01 | 76.66 | 79.90 |
+
+FP32
+
+| Batch size | Throughput Avg [mol/ms] | Latency Avg [ms] | Latency 90% [ms] |Latency 95% [ms] |Latency 99% [ms] |
+|:------------:|:------:|:-----:|:-----:|:-----:|:-----:|
+| 1600 | 3.39 | 475.86 | 473.82 | 475.64 | 891.18 |
+| 800 | 3.36 | 239.17 | 240.64 | 241.65 | 243.70 |
+| 400 | 3.17 | 126.67 | 128.19 | 128.82 | 130.54 |
+
+
+To achieve these same results, follow the steps in the [Quick Start Guide](#quick-start-guide).
+
+
+## Release notes
+
+### Changelog
+
+August 2021
+- Initial release
+
+### Known issues
+
+If you encounter `OSError: [Errno 12] Cannot allocate memory` during the Dataloader iterator creation (more precisely during the `fork()`, this is most likely due to the use of the `--precompute_bases` flag. If you cannot add more RAM or Swap to your machine, it is recommended to turn off bases precomputation by removing the `--precompute_bases` flag or using `--precompute_bases false`.

env/SE3Transformer/build/lib/se3_transformer/data_loading/__init__.py

@@ -0,0 +1 @@
+from .qm9 import QM9DataModule

env/SE3Transformer/build/lib/se3_transformer/data_loading/data_module.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import torch.distributed as dist
+from abc import ABC
+from torch.utils.data import DataLoader, DistributedSampler, Dataset
+
+from se3_transformer.runtime.utils import get_local_rank
+
+
+def _get_dataloader(dataset: Dataset, shuffle: bool, **kwargs) -> DataLoader:
+    # Classic or distributed dataloader depending on the context
+    sampler = DistributedSampler(dataset, shuffle=shuffle) if dist.is_initialized() else None
+    return DataLoader(dataset, shuffle=(shuffle and sampler is None), sampler=sampler, **kwargs)
+
+
+class DataModule(ABC):
+    """ Abstract DataModule. Children must define self.ds_{train | val | test}. """
+
+    def __init__(self, **dataloader_kwargs):
+        super().__init__()
+        if get_local_rank() == 0:
+            self.prepare_data()
+
+        # Wait until rank zero has prepared the data (download, preprocessing, ...)
+        if dist.is_initialized():
+            dist.barrier(device_ids=[get_local_rank()])
+
+        self.dataloader_kwargs = {'pin_memory': True, 'persistent_workers': True, **dataloader_kwargs}
+        self.ds_train, self.ds_val, self.ds_test = None, None, None
+
+    def prepare_data(self):
+        """ Method called only once per node. Put here any downloading or preprocessing """
+        pass
+
+    def train_dataloader(self) -> DataLoader:
+        return _get_dataloader(self.ds_train, shuffle=True, **self.dataloader_kwargs)
+
+    def val_dataloader(self) -> DataLoader:
+        return _get_dataloader(self.ds_val, shuffle=False, **self.dataloader_kwargs)
+
+    def test_dataloader(self) -> DataLoader:
+        return _get_dataloader(self.ds_test, shuffle=False, **self.dataloader_kwargs)

env/SE3Transformer/build/lib/se3_transformer/data_loading/qm9.py

@@ -0,0 +1,173 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+from typing import Tuple
+
+import dgl
+import pathlib
+import torch
+from dgl.data import QM9EdgeDataset
+from dgl import DGLGraph
+from torch import Tensor
+from torch.utils.data import random_split, DataLoader, Dataset
+from tqdm import tqdm
+
+from se3_transformer.data_loading.data_module import DataModule
+from se3_transformer.model.basis import get_basis
+from se3_transformer.runtime.utils import get_local_rank, str2bool, using_tensor_cores
+
+
+def _get_relative_pos(qm9_graph: DGLGraph) -> Tensor:
+    x = qm9_graph.ndata['pos']
+    src, dst = qm9_graph.edges()
+    rel_pos = x[dst] - x[src]
+    return rel_pos
+
+
+def _get_split_sizes(full_dataset: Dataset) -> Tuple[int, int, int]:
+    len_full = len(full_dataset)
+    len_train = 100_000
+    len_test = int(0.1 * len_full)
+    len_val = len_full - len_train - len_test
+    return len_train, len_val, len_test
+
+
+class QM9DataModule(DataModule):
+    """
+    Datamodule wrapping https://docs.dgl.ai/en/latest/api/python/dgl.data.html#qm9edge-dataset
+    Training set is 100k molecules. Test set is 10% of the dataset. Validation set is the rest.
+    This includes all the molecules from QM9 except the ones that are uncharacterized.
+    """
+
+    NODE_FEATURE_DIM = 6
+    EDGE_FEATURE_DIM = 4
+
+    def __init__(self,
+                 data_dir: pathlib.Path,
+                 task: str = 'homo',
+                 batch_size: int = 240,
+                 num_workers: int = 8,
+                 num_degrees: int = 4,
+                 amp: bool = False,
+                 precompute_bases: bool = False,
+                 **kwargs):
+        self.data_dir = data_dir  # This needs to be before __init__ so that prepare_data has access to it
+        super().__init__(batch_size=batch_size, num_workers=num_workers, collate_fn=self._collate)
+        self.amp = amp
+        self.task = task
+        self.batch_size = batch_size
+        self.num_degrees = num_degrees
+
+        qm9_kwargs = dict(label_keys=[self.task], verbose=False, raw_dir=str(data_dir))
+        if precompute_bases:
+            bases_kwargs = dict(max_degree=num_degrees - 1, use_pad_trick=using_tensor_cores(amp), amp=amp)
+            full_dataset = CachedBasesQM9EdgeDataset(bases_kwargs=bases_kwargs, batch_size=batch_size,
+                                                     num_workers=num_workers, **qm9_kwargs)
+        else:
+            full_dataset = QM9EdgeDataset(**qm9_kwargs)
+
+        self.ds_train, self.ds_val, self.ds_test = random_split(full_dataset, _get_split_sizes(full_dataset),
+                                                                generator=torch.Generator().manual_seed(0))
+
+        train_targets = full_dataset.targets[self.ds_train.indices, full_dataset.label_keys[0]]
+        self.targets_mean = train_targets.mean()
+        self.targets_std = train_targets.std()
+
+    def prepare_data(self):
+        # Download the QM9 preprocessed data
+        QM9EdgeDataset(verbose=True, raw_dir=str(self.data_dir))
+
+    def _collate(self, samples):
+        graphs, y, *bases = map(list, zip(*samples))
+        batched_graph = dgl.batch(graphs)
+        edge_feats = {'0': batched_graph.edata['edge_attr'][..., None]}
+        batched_graph.edata['rel_pos'] = _get_relative_pos(batched_graph)
+        # get node features
+        node_feats = {'0': batched_graph.ndata['attr'][:, :6, None]}
+        targets = (torch.cat(y) - self.targets_mean) / self.targets_std
+
+        if bases:
+            # collate bases
+            all_bases = {
+                key: torch.cat([b[key] for b in bases[0]], dim=0)
+                for key in bases[0][0].keys()
+            }
+
+            return batched_graph, node_feats, edge_feats, all_bases, targets
+        else:
+            return batched_graph, node_feats, edge_feats, targets
+
+    @staticmethod
+    def add_argparse_args(parent_parser):
+        parser = parent_parser.add_argument_group("QM9 dataset")
+        parser.add_argument('--task', type=str, default='homo', const='homo', nargs='?',
+                            choices=['mu', 'alpha', 'homo', 'lumo', 'gap', 'r2', 'zpve', 'U0', 'U', 'H', 'G', 'Cv',
+                                     'U0_atom', 'U_atom', 'H_atom', 'G_atom', 'A', 'B', 'C'],
+                            help='Regression task to train on')
+        parser.add_argument('--precompute_bases', type=str2bool, nargs='?', const=True, default=False,
+                            help='Precompute bases at the beginning of the script during dataset initialization,'
+                                 ' instead of computing them at the beginning of each forward pass.')
+        return parent_parser
+
+    def __repr__(self):
+        return f'QM9({self.task})'
+
+
+class CachedBasesQM9EdgeDataset(QM9EdgeDataset):
+    """ Dataset extending the QM9 dataset from DGL with precomputed (cached in RAM) pairwise bases """
+
+    def __init__(self, bases_kwargs: dict, batch_size: int, num_workers: int, *args, **kwargs):
+        """
+        :param bases_kwargs:  Arguments to feed the bases computation function
+        :param batch_size:    Batch size to use when iterating over the dataset for computing bases
+        """
+        self.bases_kwargs = bases_kwargs
+        self.batch_size = batch_size
+        self.bases = None
+        self.num_workers = num_workers
+        super().__init__(*args, **kwargs)
+
+    def load(self):
+        super().load()
+        # Iterate through the dataset and compute bases (pairwise only)
+        # Potential improvement: use multi-GPU and gather
+        dataloader = DataLoader(self, shuffle=False, batch_size=self.batch_size, num_workers=self.num_workers,
+                                collate_fn=lambda samples: dgl.batch([sample[0] for sample in samples]))
+        bases = []
+        for i, graph in tqdm(enumerate(dataloader), total=len(dataloader), desc='Precomputing QM9 bases',
+                             disable=get_local_rank() != 0):
+            rel_pos = _get_relative_pos(graph)
+            # Compute the bases with the GPU but convert the result to CPU to store in RAM
+            bases.append({k: v.cpu() for k, v in get_basis(rel_pos.cuda(), **self.bases_kwargs).items()})
+        self.bases = bases  # Assign at the end so that __getitem__ isn't confused
+
+    def __getitem__(self, idx: int):
+        graph, label = super().__getitem__(idx)
+
+        if self.bases:
+            bases_idx = idx // self.batch_size
+            bases_cumsum_idx = self.ne_cumsum[idx] - self.ne_cumsum[bases_idx * self.batch_size]
+            bases_cumsum_next_idx = self.ne_cumsum[idx + 1] - self.ne_cumsum[bases_idx * self.batch_size]
+            return graph, label, {key: basis[bases_cumsum_idx:bases_cumsum_next_idx] for key, basis in
+                                  self.bases[bases_idx].items()}
+        else:
+            return graph, label

env/SE3Transformer/build/lib/se3_transformer/model/__init__.py

@@ -0,0 +1,2 @@
+from .transformer import SE3Transformer, SE3TransformerPooled
+from .fiber import Fiber

env/SE3Transformer/build/lib/se3_transformer/model/basis.py

@@ -0,0 +1,178 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+
+from functools import lru_cache
+from typing import Dict, List
+
+import e3nn.o3 as o3
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.cuda.nvtx import range as nvtx_range
+
+from se3_transformer.runtime.utils import degree_to_dim
+
+
+@lru_cache(maxsize=None)
+def get_clebsch_gordon(J: int, d_in: int, d_out: int, device) -> Tensor:
+    """ Get the (cached) Q^{d_out,d_in}_J matrices from equation (8) """
+    return o3.wigner_3j(J, d_in, d_out, dtype=torch.float64, device=device).permute(2, 1, 0)
+
+
+@lru_cache(maxsize=None)
+def get_all_clebsch_gordon(max_degree: int, device) -> List[List[Tensor]]:
+    all_cb = []
+    for d_in in range(max_degree + 1):
+        for d_out in range(max_degree + 1):
+            K_Js = []
+            for J in range(abs(d_in - d_out), d_in + d_out + 1):
+                K_Js.append(get_clebsch_gordon(J, d_in, d_out, device))
+            all_cb.append(K_Js)
+    return all_cb
+
+
+def get_spherical_harmonics(relative_pos: Tensor, max_degree: int) -> List[Tensor]:
+    all_degrees = list(range(2 * max_degree + 1))
+    with nvtx_range('spherical harmonics'):
+        sh = o3.spherical_harmonics(all_degrees, relative_pos, normalize=True)
+        return torch.split(sh, [degree_to_dim(d) for d in all_degrees], dim=1)
+
+
+@torch.jit.script
+def get_basis_script(max_degree: int,
+                     use_pad_trick: bool,
+                     spherical_harmonics: List[Tensor],
+                     clebsch_gordon: List[List[Tensor]],
+                     amp: bool) -> Dict[str, Tensor]:
+    """
+    Compute pairwise bases matrices for degrees up to max_degree
+    :param max_degree:            Maximum input or output degree
+    :param use_pad_trick:         Pad some of the odd dimensions for a better use of Tensor Cores
+    :param spherical_harmonics:   List of computed spherical harmonics
+    :param clebsch_gordon:        List of computed CB-coefficients
+    :param amp:                   When true, return bases in FP16 precision
+    """
+    basis = {}
+    idx = 0
+    # Double for loop instead of product() because of JIT script
+    for d_in in range(max_degree + 1):
+        for d_out in range(max_degree + 1):
+            key = f'{d_in},{d_out}'
+            K_Js = []
+            for freq_idx, J in enumerate(range(abs(d_in - d_out), d_in + d_out + 1)):
+                Q_J = clebsch_gordon[idx][freq_idx]
+                K_Js.append(torch.einsum('n f, k l f -> n l k', spherical_harmonics[J].float(), Q_J.float()))
+
+            basis[key] = torch.stack(K_Js, 2)  # Stack on second dim so order is n l f k
+            if amp:
+                basis[key] = basis[key].half()
+            if use_pad_trick:
+                basis[key] = F.pad(basis[key], (0, 1))  # Pad the k dimension, that can be sliced later
+
+            idx += 1
+
+    return basis
+
+
+@torch.jit.script
+def update_basis_with_fused(basis: Dict[str, Tensor],
+                            max_degree: int,
+                            use_pad_trick: bool,
+                            fully_fused: bool) -> Dict[str, Tensor]:
+    """ Update the basis dict with partially and optionally fully fused bases """
+    num_edges = basis['0,0'].shape[0]
+    device = basis['0,0'].device
+    dtype = basis['0,0'].dtype
+    sum_dim = sum([degree_to_dim(d) for d in range(max_degree + 1)])
+
+    # Fused per output degree
+    for d_out in range(max_degree + 1):
+        sum_freq = sum([degree_to_dim(min(d, d_out)) for d in range(max_degree + 1)])
+        basis_fused = torch.zeros(num_edges, sum_dim, sum_freq, degree_to_dim(d_out) + int(use_pad_trick),
+                                  device=device, dtype=dtype)
+        acc_d, acc_f = 0, 0
+        for d_in in range(max_degree + 1):
+            basis_fused[:, acc_d:acc_d + degree_to_dim(d_in), acc_f:acc_f + degree_to_dim(min(d_out, d_in)),
+            :degree_to_dim(d_out)] = basis[f'{d_in},{d_out}'][:, :, :, :degree_to_dim(d_out)]
+
+            acc_d += degree_to_dim(d_in)
+            acc_f += degree_to_dim(min(d_out, d_in))
+
+        basis[f'out{d_out}_fused'] = basis_fused
+
+    # Fused per input degree
+    for d_in in range(max_degree + 1):
+        sum_freq = sum([degree_to_dim(min(d, d_in)) for d in range(max_degree + 1)])
+        basis_fused = torch.zeros(num_edges, degree_to_dim(d_in), sum_freq, sum_dim,
+                                  device=device, dtype=dtype)
+        acc_d, acc_f = 0, 0
+        for d_out in range(max_degree + 1):
+            basis_fused[:, :, acc_f:acc_f + degree_to_dim(min(d_out, d_in)), acc_d:acc_d + degree_to_dim(d_out)] \
+                = basis[f'{d_in},{d_out}'][:, :, :, :degree_to_dim(d_out)]
+
+            acc_d += degree_to_dim(d_out)
+            acc_f += degree_to_dim(min(d_out, d_in))
+
+        basis[f'in{d_in}_fused'] = basis_fused
+
+    if fully_fused:
+        # Fully fused
+        # Double sum this way because of JIT script
+        sum_freq = sum([
+            sum([degree_to_dim(min(d_in, d_out)) for d_in in range(max_degree + 1)]) for d_out in range(max_degree + 1)
+        ])
+        basis_fused = torch.zeros(num_edges, sum_dim, sum_freq, sum_dim, device=device, dtype=dtype)
+
+        acc_d, acc_f = 0, 0
+        for d_out in range(max_degree + 1):
+            b = basis[f'out{d_out}_fused']
+            basis_fused[:, :, acc_f:acc_f + b.shape[2], acc_d:acc_d + degree_to_dim(d_out)] = b[:, :, :,
+                                                                                              :degree_to_dim(d_out)]
+            acc_f += b.shape[2]
+            acc_d += degree_to_dim(d_out)
+
+        basis['fully_fused'] = basis_fused
+
+    del basis['0,0']  # We know that the basis for l = k = 0 is filled with a constant
+    return basis
+
+
+def get_basis(relative_pos: Tensor,
+              max_degree: int = 4,
+              compute_gradients: bool = False,
+              use_pad_trick: bool = False,
+              amp: bool = False) -> Dict[str, Tensor]:
+    with nvtx_range('spherical harmonics'):
+        spherical_harmonics = get_spherical_harmonics(relative_pos, max_degree)
+    with nvtx_range('CB coefficients'):
+        clebsch_gordon = get_all_clebsch_gordon(max_degree, relative_pos.device)
+
+    with torch.autograd.set_grad_enabled(compute_gradients):
+        with nvtx_range('bases'):
+            basis = get_basis_script(max_degree=max_degree,
+                                     use_pad_trick=use_pad_trick,
+                                     spherical_harmonics=spherical_harmonics,
+                                     clebsch_gordon=clebsch_gordon,
+                                     amp=amp)
+            return basis

env/SE3Transformer/build/lib/se3_transformer/model/fiber.py

@@ -0,0 +1,144 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+
+from collections import namedtuple
+from itertools import product
+from typing import Dict
+
+import torch
+from torch import Tensor
+
+from se3_transformer.runtime.utils import degree_to_dim
+
+FiberEl = namedtuple('FiberEl', ['degree', 'channels'])
+
+
+class Fiber(dict):
+    """
+    Describes the structure of some set of features.
+    Features are split into types (0, 1, 2, 3, ...). A feature of type k has a dimension of 2k+1.
+    Type-0 features: invariant scalars
+    Type-1 features: equivariant 3D vectors
+    Type-2 features: equivariant symmetric traceless matrices
+    ...
+
+    As inputs to a SE3 layer, there can be many features of the same types, and many features of different types.
+    The 'multiplicity' or 'number of channels' is the number of features of a given type.
+    This class puts together all the degrees and their multiplicities in order to describe
+        the inputs, outputs or hidden features of SE3 layers.
+    """
+
+    def __init__(self, structure):
+        if isinstance(structure, dict):
+            structure = [FiberEl(int(d), int(m)) for d, m in sorted(structure.items(), key=lambda x: x[1])]
+        elif not isinstance(structure[0], FiberEl):
+            structure = list(map(lambda t: FiberEl(*t), sorted(structure, key=lambda x: x[1])))
+        self.structure = structure
+        super().__init__({d: m for d, m in self.structure})
+
+    @property
+    def degrees(self):
+        return sorted([t.degree for t in self.structure])
+
+    @property
+    def channels(self):
+        return [self[d] for d in self.degrees]
+
+    @property
+    def num_features(self):
+        """ Size of the resulting tensor if all features were concatenated together """
+        return sum(t.channels * degree_to_dim(t.degree) for t in self.structure)
+
+    @staticmethod
+    def create(num_degrees: int, num_channels: int):
+        """ Create a Fiber with degrees 0..num_degrees-1, all with the same multiplicity """
+        return Fiber([(degree, num_channels) for degree in range(num_degrees)])
+
+    @staticmethod
+    def from_features(feats: Dict[str, Tensor]):
+        """ Infer the Fiber structure from a feature dict """
+        structure = {}
+        for k, v in feats.items():
+            degree = int(k)
+            assert len(v.shape) == 3, 'Feature shape should be (N, C, 2D+1)'
+            assert v.shape[-1] == degree_to_dim(degree)
+            structure[degree] = v.shape[-2]
+        return Fiber(structure)
+
+    def __getitem__(self, degree: int):
+        """ fiber[degree] returns the multiplicity for this degree """
+        return dict(self.structure).get(degree, 0)
+
+    def __iter__(self):
+        """ Iterate over namedtuples (degree, channels) """
+        return iter(self.structure)
+
+    def __mul__(self, other):
+        """
+        If other in an int, multiplies all the multiplicities by other.
+        If other is a fiber, returns the cartesian product.
+        """
+        if isinstance(other, Fiber):
+            return product(self.structure, other.structure)
+        elif isinstance(other, int):
+            return Fiber({t.degree: t.channels * other for t in self.structure})
+
+    def __add__(self, other):
+        """
+        If other in an int, add other to all the multiplicities.
+        If other is a fiber, add the multiplicities of the fibers together.
+        """
+        if isinstance(other, Fiber):
+            return Fiber({t.degree: t.channels + other[t.degree] for t in self.structure})
+        elif isinstance(other, int):
+            return Fiber({t.degree: t.channels + other for t in self.structure})
+
+    def __repr__(self):
+        return str(self.structure)
+
+    @staticmethod
+    def combine_max(f1, f2):
+        """ Combine two fiber by taking the maximum multiplicity for each degree in both fibers """
+        new_dict = dict(f1.structure)
+        for k, m in f2.structure:
+            new_dict[k] = max(new_dict.get(k, 0), m)
+
+        return Fiber(list(new_dict.items()))
+
+    @staticmethod
+    def combine_selectively(f1, f2):
+        """ Combine two fiber by taking the sum of multiplicities for each degree in the first fiber """
+        # only use orders which occur in fiber f1
+        new_dict = dict(f1.structure)
+        for k in f1.degrees:
+            if k in f2.degrees:
+                new_dict[k] += f2[k]
+        return Fiber(list(new_dict.items()))
+
+    def to_attention_heads(self, tensors: Dict[str, Tensor], num_heads: int):
+        # dict(N, num_channels, 2d+1) -> (N, num_heads, -1)
+        fibers = [tensors[str(degree)].reshape(*tensors[str(degree)].shape[:-2], num_heads, -1) for degree in
+                  self.degrees]
+        fibers = torch.cat(fibers, -1)
+        return fibers

env/SE3Transformer/build/lib/se3_transformer/model/layers/__init__.py

@@ -0,0 +1,5 @@
+from .linear import LinearSE3
+from .norm import NormSE3
+from .pooling import GPooling
+from .convolution import ConvSE3
+from .attention import AttentionBlockSE3

env/SE3Transformer/build/lib/se3_transformer/model/layers/attention.py

@@ -0,0 +1,180 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import dgl
+import numpy as np
+import torch
+import torch.nn as nn
+from dgl import DGLGraph
+from dgl.ops import edge_softmax
+from torch import Tensor
+from typing import Dict, Optional, Union
+
+from se3_transformer.model.fiber import Fiber
+from se3_transformer.model.layers.convolution import ConvSE3, ConvSE3FuseLevel
+from se3_transformer.model.layers.linear import LinearSE3
+from se3_transformer.runtime.utils import degree_to_dim, aggregate_residual, unfuse_features
+from torch.cuda.nvtx import range as nvtx_range
+
+
+class AttentionSE3(nn.Module):
+    """ Multi-headed sparse graph self-attention (SE(3)-equivariant) """
+
+    def __init__(
+            self,
+            num_heads: int,
+            key_fiber: Fiber,
+            value_fiber: Fiber
+    ):
+        """
+        :param num_heads:     Number of attention heads
+        :param key_fiber:     Fiber for the keys (and also for the queries)
+        :param value_fiber:   Fiber for the values
+        """
+        super().__init__()
+        self.num_heads = num_heads
+        self.key_fiber = key_fiber
+        self.value_fiber = value_fiber
+
+    def forward(
+            self,
+            value: Union[Tensor, Dict[str, Tensor]],  # edge features (may be fused)
+            key: Union[Tensor, Dict[str, Tensor]],  # edge features (may be fused)
+            query: Dict[str, Tensor],  # node features
+            graph: DGLGraph
+    ):
+        with nvtx_range('AttentionSE3'):
+            with nvtx_range('reshape keys and queries'):
+                if isinstance(key, Tensor):
+                    # case where features of all types are fused
+                    key = key.reshape(key.shape[0], self.num_heads, -1)
+                    # need to reshape queries that way to keep the same layout as keys
+                    out = torch.cat([query[str(d)] for d in self.key_fiber.degrees], dim=-1)
+                    query = out.reshape(list(query.values())[0].shape[0], self.num_heads, -1)
+                else:
+                    # features are not fused, need to fuse and reshape them
+                    key = self.key_fiber.to_attention_heads(key, self.num_heads)
+                    query = self.key_fiber.to_attention_heads(query, self.num_heads)
+
+            with nvtx_range('attention dot product + softmax'):
+                # Compute attention weights (softmax of inner product between key and query)
+                edge_weights = dgl.ops.e_dot_v(graph, key, query).squeeze(-1)
+                edge_weights /= np.sqrt(self.key_fiber.num_features)
+                edge_weights = edge_softmax(graph, edge_weights)
+                edge_weights = edge_weights[..., None, None]
+
+            with nvtx_range('weighted sum'):
+                if isinstance(value, Tensor):
+                    # features of all types are fused
+                    v = value.view(value.shape[0], self.num_heads, -1, value.shape[-1])
+                    weights = edge_weights * v
+                    feat_out = dgl.ops.copy_e_sum(graph, weights)
+                    feat_out = feat_out.view(feat_out.shape[0], -1, feat_out.shape[-1])  # merge heads
+                    out = unfuse_features(feat_out, self.value_fiber.degrees)
+                else:
+                    out = {}
+                    for degree, channels in self.value_fiber:
+                        v = value[str(degree)].view(-1, self.num_heads, channels // self.num_heads,
+                                                    degree_to_dim(degree))
+                        weights = edge_weights * v
+                        res = dgl.ops.copy_e_sum(graph, weights)
+                        out[str(degree)] = res.view(-1, channels, degree_to_dim(degree))  # merge heads
+
+                return out
+
+
+class AttentionBlockSE3(nn.Module):
+    """ Multi-headed sparse graph self-attention block with skip connection, linear projection (SE(3)-equivariant) """
+
+    def __init__(
+            self,
+            fiber_in: Fiber,
+            fiber_out: Fiber,
+            fiber_edge: Optional[Fiber] = None,
+            num_heads: int = 4,
+            channels_div: int = 2,
+            use_layer_norm: bool = False,
+            max_degree: bool = 4,
+            fuse_level: ConvSE3FuseLevel = ConvSE3FuseLevel.FULL,
+            **kwargs
+    ):
+        """
+        :param fiber_in:         Fiber describing the input features
+        :param fiber_out:        Fiber describing the output features
+        :param fiber_edge:       Fiber describing the edge features (node distances excluded)
+        :param num_heads:        Number of attention heads
+        :param channels_div:     Divide the channels by this integer for computing values
+        :param use_layer_norm:   Apply layer normalization between MLP layers
+        :param max_degree:       Maximum degree used in the bases computation
+        :param fuse_level:       Maximum fuse level to use in TFN convolutions
+        """
+        super().__init__()
+        if fiber_edge is None:
+            fiber_edge = Fiber({})
+        self.fiber_in = fiber_in
+        # value_fiber has same structure as fiber_out but #channels divided by 'channels_div'
+        value_fiber = Fiber([(degree, channels // channels_div) for degree, channels in fiber_out])
+        # key_query_fiber has the same structure as fiber_out, but only degrees which are in in_fiber
+        # (queries are merely projected, hence degrees have to match input)
+        key_query_fiber = Fiber([(fe.degree, fe.channels) for fe in value_fiber if fe.degree in fiber_in.degrees])
+
+        self.to_key_value = ConvSE3(fiber_in, value_fiber + key_query_fiber, pool=False, fiber_edge=fiber_edge,
+                                    use_layer_norm=use_layer_norm, max_degree=max_degree, fuse_level=fuse_level,
+                                    allow_fused_output=True)
+        self.to_query = LinearSE3(fiber_in, key_query_fiber)
+        self.attention = AttentionSE3(num_heads, key_query_fiber, value_fiber)
+        self.project = LinearSE3(value_fiber + fiber_in, fiber_out)
+
+    def forward(
+            self,
+            node_features: Dict[str, Tensor],
+            edge_features: Dict[str, Tensor],
+            graph: DGLGraph,
+            basis: Dict[str, Tensor]
+    ):
+        with nvtx_range('AttentionBlockSE3'):
+            with nvtx_range('keys / values'):
+                fused_key_value = self.to_key_value(node_features, edge_features, graph, basis)
+                key, value = self._get_key_value_from_fused(fused_key_value)
+
+            with nvtx_range('queries'):
+                query = self.to_query(node_features)
+            
+            z = self.attention(value, key, query, graph)
+            z_concat = aggregate_residual(node_features, z, 'cat')
+            return self.project(z_concat)
+
+    def _get_key_value_from_fused(self, fused_key_value):
+        # Extract keys and queries features from fused features
+        if isinstance(fused_key_value, Tensor):
+            # Previous layer was a fully fused convolution
+            value, key = torch.chunk(fused_key_value, chunks=2, dim=-2)
+        else:
+            key, value = {}, {}
+            for degree, feat in fused_key_value.items():
+                if int(degree) in self.fiber_in.degrees:
+                    value[degree], key[degree] = torch.chunk(feat, chunks=2, dim=-2)
+                else:
+                    value[degree] = feat
+
+        return key, value

env/SE3Transformer/build/lib/se3_transformer/model/layers/convolution.py

@@ -0,0 +1,336 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+from enum import Enum
+from itertools import product
+from typing import Dict
+
+import dgl
+import numpy as np
+import torch
+import torch.nn as nn
+from dgl import DGLGraph
+from torch import Tensor
+from torch.cuda.nvtx import range as nvtx_range
+
+from se3_transformer.model.fiber import Fiber
+from se3_transformer.runtime.utils import degree_to_dim, unfuse_features
+
+
+class ConvSE3FuseLevel(Enum):
+    """
+    Enum to select a maximum level of fusing optimizations that will be applied when certain conditions are met.
+    If a desired level L is picked and the level L cannot be applied to a level, other fused ops < L are considered.
+    A higher level means faster training, but also more memory usage.
+    If you are tight on memory and want to feed large inputs to the network, choose a low value.
+    If you want to train fast, choose a high value.
+    Recommended value is FULL with AMP.
+
+    Fully fused TFN convolutions requirements:
+    - all input channels are the same
+    - all output channels are the same
+    - input degrees span the range [0, ..., max_degree]
+    - output degrees span the range [0, ..., max_degree]
+
+    Partially fused TFN convolutions requirements:
+    * For fusing by output degree:
+    - all input channels are the same
+    - input degrees span the range [0, ..., max_degree]
+    * For fusing by input degree:
+    - all output channels are the same
+    - output degrees span the range [0, ..., max_degree]
+
+    Original TFN pairwise convolutions: no requirements
+    """
+
+    FULL = 2
+    PARTIAL = 1
+    NONE = 0
+
+
+class RadialProfile(nn.Module):
+    """
+    Radial profile function.
+    Outputs weights used to weigh basis matrices in order to get convolution kernels.
+    In TFN notation: $R^{l,k}$
+    In SE(3)-Transformer notation: $\phi^{l,k}$
+
+    Note:
+        In the original papers, this function only depends on relative node distances ||x||.
+        Here, we allow this function to also take as input additional invariant edge features.
+        This does not break equivariance and adds expressive power to the model.
+
+    Diagram:
+        invariant edge features (node distances included) ───> MLP layer (shared across edges) ───> radial weights
+    """
+
+    def __init__(
+            self,
+            num_freq: int,
+            channels_in: int,
+            channels_out: int,
+            edge_dim: int = 1,
+            mid_dim: int = 32,
+            use_layer_norm: bool = False
+    ):
+        """
+        :param num_freq:         Number of frequencies
+        :param channels_in:      Number of input channels
+        :param channels_out:     Number of output channels
+        :param edge_dim:         Number of invariant edge features (input to the radial function)
+        :param mid_dim:          Size of the hidden MLP layers
+        :param use_layer_norm:   Apply layer normalization between MLP layers
+        """
+        super().__init__()
+        modules = [
+            nn.Linear(edge_dim, mid_dim),
+            nn.LayerNorm(mid_dim) if use_layer_norm else None,
+            nn.ReLU(),
+            nn.Linear(mid_dim, mid_dim),
+            nn.LayerNorm(mid_dim) if use_layer_norm else None,
+            nn.ReLU(),
+            nn.Linear(mid_dim, num_freq * channels_in * channels_out, bias=False)
+        ]
+
+        self.net = nn.Sequential(*[m for m in modules if m is not None])
+
+    def forward(self, features: Tensor) -> Tensor:
+        return self.net(features)
+
+
+class VersatileConvSE3(nn.Module):
+    """
+    Building block for TFN convolutions.
+    This single module can be used for fully fused convolutions, partially fused convolutions, or pairwise convolutions.
+    """
+
+    def __init__(self,
+                 freq_sum: int,
+                 channels_in: int,
+                 channels_out: int,
+                 edge_dim: int,
+                 use_layer_norm: bool,
+                 fuse_level: ConvSE3FuseLevel):
+        super().__init__()
+        self.freq_sum = freq_sum
+        self.channels_out = channels_out
+        self.channels_in = channels_in
+        self.fuse_level = fuse_level
+        self.radial_func = RadialProfile(num_freq=freq_sum,
+                                         channels_in=channels_in,
+                                         channels_out=channels_out,
+                                         edge_dim=edge_dim,
+                                         use_layer_norm=use_layer_norm)
+
+    def forward(self, features: Tensor, invariant_edge_feats: Tensor, basis: Tensor):
+        with nvtx_range(f'VersatileConvSE3'):
+            num_edges = features.shape[0]
+            in_dim = features.shape[2]
+            with nvtx_range(f'RadialProfile'):
+                radial_weights = self.radial_func(invariant_edge_feats) \
+                    .view(-1, self.channels_out, self.channels_in * self.freq_sum)
+
+            if basis is not None:
+                # This block performs the einsum n i l, n o i f, n l f k -> n o k
+                out_dim = basis.shape[-1]
+                if self.fuse_level != ConvSE3FuseLevel.FULL:
+                    out_dim += out_dim % 2 - 1  # Account for padded basis
+                basis_view = basis.view(num_edges, in_dim, -1)
+                tmp = (features @ basis_view).view(num_edges, -1, basis.shape[-1])
+                return (radial_weights @ tmp)[:, :, :out_dim]
+            else:
+                # k = l = 0 non-fused case
+                return radial_weights @ features
+
+
+class ConvSE3(nn.Module):
+    """
+    SE(3)-equivariant graph convolution (Tensor Field Network convolution).
+    This convolution can map an arbitrary input Fiber to an arbitrary output Fiber, while preserving equivariance.
+    Features of different degrees interact together to produce output features.
+
+    Note 1:
+        The option is given to not pool the output. This means that the convolution sum over neighbors will not be
+        done, and the returned features will be edge features instead of node features.
+
+    Note 2:
+        Unlike the original paper and implementation, this convolution can handle edge feature of degree greater than 0.
+        Input edge features are concatenated with input source node features before the kernel is applied.
+     """
+
+    def __init__(
+            self,
+            fiber_in: Fiber,
+            fiber_out: Fiber,
+            fiber_edge: Fiber,
+            pool: bool = True,
+            use_layer_norm: bool = False,
+            self_interaction: bool = False,
+            max_degree: int = 4,
+            fuse_level: ConvSE3FuseLevel = ConvSE3FuseLevel.FULL,
+            allow_fused_output: bool = False
+    ):
+        """
+        :param fiber_in:           Fiber describing the input features
+        :param fiber_out:          Fiber describing the output features
+        :param fiber_edge:         Fiber describing the edge features (node distances excluded)
+        :param pool:               If True, compute final node features by averaging incoming edge features
+        :param use_layer_norm:     Apply layer normalization between MLP layers
+        :param self_interaction:   Apply self-interaction of nodes
+        :param max_degree:         Maximum degree used in the bases computation
+        :param fuse_level:         Maximum fuse level to use in TFN convolutions
+        :param allow_fused_output: Allow the module to output a fused representation of features
+        """
+        super().__init__()
+        self.pool = pool
+        self.fiber_in = fiber_in
+        self.fiber_out = fiber_out
+        self.self_interaction = self_interaction
+        self.max_degree = max_degree
+        self.allow_fused_output = allow_fused_output
+
+        # channels_in: account for the concatenation of edge features
+        channels_in_set = set([f.channels + fiber_edge[f.degree] * (f.degree > 0) for f in self.fiber_in])
+        channels_out_set = set([f.channels for f in self.fiber_out])
+        unique_channels_in = (len(channels_in_set) == 1)
+        unique_channels_out = (len(channels_out_set) == 1)
+        degrees_up_to_max = list(range(max_degree + 1))
+        common_args = dict(edge_dim=fiber_edge[0] + 1, use_layer_norm=use_layer_norm)
+
+        if fuse_level.value >= ConvSE3FuseLevel.FULL.value and \
+                unique_channels_in and fiber_in.degrees == degrees_up_to_max and \
+                unique_channels_out and fiber_out.degrees == degrees_up_to_max:
+            # Single fused convolution
+            self.used_fuse_level = ConvSE3FuseLevel.FULL
+
+            sum_freq = sum([
+                degree_to_dim(min(d_in, d_out))
+                for d_in, d_out in product(degrees_up_to_max, degrees_up_to_max)
+            ])
+
+            self.conv = VersatileConvSE3(sum_freq, list(channels_in_set)[0], list(channels_out_set)[0],
+                                         fuse_level=self.used_fuse_level, **common_args)
+
+        elif fuse_level.value >= ConvSE3FuseLevel.PARTIAL.value and \
+                unique_channels_in and fiber_in.degrees == degrees_up_to_max:
+            # Convolutions fused per output degree
+            self.used_fuse_level = ConvSE3FuseLevel.PARTIAL
+            self.conv_out = nn.ModuleDict()
+            for d_out, c_out in fiber_out:
+                sum_freq = sum([degree_to_dim(min(d_out, d)) for d in fiber_in.degrees])
+                self.conv_out[str(d_out)] = VersatileConvSE3(sum_freq, list(channels_in_set)[0], c_out,
+                                                             fuse_level=self.used_fuse_level, **common_args)
+
+        elif fuse_level.value >= ConvSE3FuseLevel.PARTIAL.value and \
+                unique_channels_out and fiber_out.degrees == degrees_up_to_max:
+            # Convolutions fused per input degree
+            self.used_fuse_level = ConvSE3FuseLevel.PARTIAL
+            self.conv_in = nn.ModuleDict()
+            for d_in, c_in in fiber_in:
+                sum_freq = sum([degree_to_dim(min(d_in, d)) for d in fiber_out.degrees])
+                self.conv_in[str(d_in)] = VersatileConvSE3(sum_freq, c_in, list(channels_out_set)[0],
+                                                           fuse_level=ConvSE3FuseLevel.FULL, **common_args)
+                                                           #fuse_level=self.used_fuse_level, **common_args)
+        else:
+            # Use pairwise TFN convolutions
+            self.used_fuse_level = ConvSE3FuseLevel.NONE
+            self.conv = nn.ModuleDict()
+            for (degree_in, channels_in), (degree_out, channels_out) in (self.fiber_in * self.fiber_out):
+                dict_key = f'{degree_in},{degree_out}'
+                channels_in_new = channels_in + fiber_edge[degree_in] * (degree_in > 0)
+                sum_freq = degree_to_dim(min(degree_in, degree_out))
+                self.conv[dict_key] = VersatileConvSE3(sum_freq, channels_in_new, channels_out,
+                                                       fuse_level=self.used_fuse_level, **common_args)
+
+        if self_interaction:
+            self.to_kernel_self = nn.ParameterDict()
+            for degree_out, channels_out in fiber_out:
+                if fiber_in[degree_out]:
+                    self.to_kernel_self[str(degree_out)] = nn.Parameter(
+                        torch.randn(channels_out, fiber_in[degree_out]) / np.sqrt(fiber_in[degree_out]))
+
+    def forward(
+            self,
+            node_feats: Dict[str, Tensor],
+            edge_feats: Dict[str, Tensor],
+            graph: DGLGraph,
+            basis: Dict[str, Tensor]
+    ):
+        with nvtx_range(f'ConvSE3'):
+            invariant_edge_feats = edge_feats['0'].squeeze(-1)
+            src, dst = graph.edges()
+            out = {}
+            in_features = []
+
+            # Fetch all input features from edge and node features
+            for degree_in in self.fiber_in.degrees:
+                src_node_features = node_feats[str(degree_in)][src]
+                if degree_in > 0 and str(degree_in) in edge_feats:
+                    # Handle edge features of any type by concatenating them to node features
+                    src_node_features = torch.cat([src_node_features, edge_feats[str(degree_in)]], dim=1)
+                in_features.append(src_node_features)
+
+            if self.used_fuse_level == ConvSE3FuseLevel.FULL:
+                in_features_fused = torch.cat(in_features, dim=-1)
+                out = self.conv(in_features_fused, invariant_edge_feats, basis['fully_fused'])
+
+                if not self.allow_fused_output or self.self_interaction or self.pool:
+                    out = unfuse_features(out, self.fiber_out.degrees)
+
+            elif self.used_fuse_level == ConvSE3FuseLevel.PARTIAL and hasattr(self, 'conv_out'):
+                in_features_fused = torch.cat(in_features, dim=-1)
+                for degree_out in self.fiber_out.degrees:
+                    out[str(degree_out)] = self.conv_out[str(degree_out)](in_features_fused, invariant_edge_feats,
+                                                                          basis[f'out{degree_out}_fused'])
+
+            elif self.used_fuse_level == ConvSE3FuseLevel.PARTIAL and hasattr(self, 'conv_in'):
+                out = 0
+                for degree_in, feature in zip(self.fiber_in.degrees, in_features):
+                    out += self.conv_in[str(degree_in)](feature, invariant_edge_feats,
+                                                        basis[f'in{degree_in}_fused'])
+                if not self.allow_fused_output or self.self_interaction or self.pool:
+                    out = unfuse_features(out, self.fiber_out.degrees)
+            else:
+                # Fallback to pairwise TFN convolutions
+                for degree_out in self.fiber_out.degrees:
+                    out_feature = 0
+                    for degree_in, feature in zip(self.fiber_in.degrees, in_features):
+                        dict_key = f'{degree_in},{degree_out}'
+                        out_feature = out_feature + self.conv[dict_key](feature, invariant_edge_feats,
+                                                                        basis.get(dict_key, None))
+                    out[str(degree_out)] = out_feature
+
+            for degree_out in self.fiber_out.degrees:
+                if self.self_interaction and str(degree_out) in self.to_kernel_self:
+                    with nvtx_range(f'self interaction'):
+                        dst_features = node_feats[str(degree_out)][dst]
+                        kernel_self = self.to_kernel_self[str(degree_out)]
+                        out[str(degree_out)] += kernel_self @ dst_features
+
+                if self.pool:
+                    with nvtx_range(f'pooling'):
+                        if isinstance(out, dict):
+                            out[str(degree_out)] = dgl.ops.copy_e_sum(graph, out[str(degree_out)])
+                        else:
+                            out = dgl.ops.copy_e_sum(graph, out)
+            return out

env/SE3Transformer/build/lib/se3_transformer/model/layers/linear.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+
+from typing import Dict
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from se3_transformer.model.fiber import Fiber
+
+
+class LinearSE3(nn.Module):
+    """
+    Graph Linear SE(3)-equivariant layer, equivalent to a 1x1 convolution.
+    Maps a fiber to a fiber with the same degrees (channels may be different).
+    No interaction between degrees, but interaction between channels.
+
+    type-0 features (C_0 channels) ────> Linear(bias=False) ────> type-0 features (C'_0 channels)
+    type-1 features (C_1 channels) ────> Linear(bias=False) ────> type-1 features (C'_1 channels)
+                                                 :
+    type-k features (C_k channels) ────> Linear(bias=False) ────> type-k features (C'_k channels)
+    """
+
+    def __init__(self, fiber_in: Fiber, fiber_out: Fiber):
+        super().__init__()
+        self.weights = nn.ParameterDict({
+            str(degree_out): nn.Parameter(
+                torch.randn(channels_out, fiber_in[degree_out]) / np.sqrt(fiber_in[degree_out]))
+            for degree_out, channels_out in fiber_out
+        })
+
+    def forward(self, features: Dict[str, Tensor], *args, **kwargs) -> Dict[str, Tensor]:
+        return {
+            degree: self.weights[degree] @ features[degree]
+            for degree, weight in self.weights.items()
+        }

env/SE3Transformer/build/lib/se3_transformer/model/layers/norm.py

@@ -0,0 +1,83 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+
+from typing import Dict
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.cuda.nvtx import range as nvtx_range
+
+from se3_transformer.model.fiber import Fiber
+
+
+class NormSE3(nn.Module):
+    """
+    Norm-based SE(3)-equivariant nonlinearity.
+
+                 ┌──> feature_norm ──> LayerNorm() ──> ReLU() ──┐
+    feature_in ──┤                                              * ──> feature_out
+                 └──> feature_phase ────────────────────────────┘
+    """
+
+    NORM_CLAMP = 2 ** -24  # Minimum positive subnormal for FP16
+
+    def __init__(self, fiber: Fiber, nonlinearity: nn.Module = nn.ReLU()):
+        super().__init__()
+        self.fiber = fiber
+        self.nonlinearity = nonlinearity
+
+        if len(set(fiber.channels)) == 1:
+            # Fuse all the layer normalizations into a group normalization
+            self.group_norm = nn.GroupNorm(num_groups=len(fiber.degrees), num_channels=sum(fiber.channels))
+        else:
+            # Use multiple layer normalizations
+            self.layer_norms = nn.ModuleDict({
+                str(degree): nn.LayerNorm(channels)
+                for degree, channels in fiber
+            })
+
+    def forward(self, features: Dict[str, Tensor], *args, **kwargs) -> Dict[str, Tensor]:
+        with nvtx_range('NormSE3'):
+            output = {}
+            if hasattr(self, 'group_norm'):
+                # Compute per-degree norms of features
+                norms = [features[str(d)].norm(dim=-1, keepdim=True).clamp(min=self.NORM_CLAMP)
+                         for d in self.fiber.degrees]
+                fused_norms = torch.cat(norms, dim=-2)
+
+                # Transform the norms only
+                new_norms = self.nonlinearity(self.group_norm(fused_norms.squeeze(-1))).unsqueeze(-1)
+                new_norms = torch.chunk(new_norms, chunks=len(self.fiber.degrees), dim=-2)
+
+                # Scale features to the new norms
+                for norm, new_norm, d in zip(norms, new_norms, self.fiber.degrees):
+                    output[str(d)] = features[str(d)] / norm * new_norm
+            else:
+                for degree, feat in features.items():
+                    norm = feat.norm(dim=-1, keepdim=True).clamp(min=self.NORM_CLAMP)
+                    new_norm = self.nonlinearity(self.layer_norms[degree](norm.squeeze(-1)).unsqueeze(-1))
+                    output[degree] = new_norm * feat / norm
+
+            return output

env/SE3Transformer/build/lib/se3_transformer/model/layers/pooling.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+from typing import Dict, Literal
+
+import torch.nn as nn
+from dgl import DGLGraph
+from dgl.nn.pytorch import AvgPooling, MaxPooling
+from torch import Tensor
+
+
+class GPooling(nn.Module):
+    """
+    Graph max/average pooling on a given feature type.
+    The average can be taken for any feature type, and equivariance will be maintained.
+    The maximum can only be taken for invariant features (type 0).
+    If you want max-pooling for type > 0 features, look into Vector Neurons.
+    """
+
+    def __init__(self, feat_type: int = 0, pool: Literal['max', 'avg'] = 'max'):
+        """
+        :param feat_type: Feature type to pool
+        :param pool: Type of pooling: max or avg
+        """
+        super().__init__()
+        assert pool in ['max', 'avg'], f'Unknown pooling: {pool}'
+        assert feat_type == 0 or pool == 'avg', 'Max pooling on type > 0 features will break equivariance'
+        self.feat_type = feat_type
+        self.pool = MaxPooling() if pool == 'max' else AvgPooling()
+
+    def forward(self, features: Dict[str, Tensor], graph: DGLGraph, **kwargs) -> Tensor:
+        pooled = self.pool(graph, features[str(self.feat_type)])
+        return pooled.squeeze(dim=-1)

env/SE3Transformer/build/lib/se3_transformer/model/transformer.py

@@ -0,0 +1,222 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import logging
+from typing import Optional, Literal, Dict
+
+import torch
+import torch.nn as nn
+from dgl import DGLGraph
+from torch import Tensor
+
+from se3_transformer.model.basis import get_basis, update_basis_with_fused
+from se3_transformer.model.layers.attention import AttentionBlockSE3
+from se3_transformer.model.layers.convolution import ConvSE3, ConvSE3FuseLevel
+from se3_transformer.model.layers.norm import NormSE3
+from se3_transformer.model.layers.pooling import GPooling
+from se3_transformer.runtime.utils import str2bool
+from se3_transformer.model.fiber import Fiber
+
+
+class Sequential(nn.Sequential):
+    """ Sequential module with arbitrary forward args and kwargs. Used to pass graph, basis and edge features. """
+
+    def forward(self, input, *args, **kwargs):
+        for module in self:
+            input = module(input, *args, **kwargs)
+        return input
+
+
+def get_populated_edge_features(relative_pos: Tensor, edge_features: Optional[Dict[str, Tensor]] = None):
+    """ Add relative positions to existing edge features """
+    edge_features = edge_features.copy() if edge_features else {}
+    r = relative_pos.norm(dim=-1, keepdim=True)
+    if '0' in edge_features:
+        edge_features['0'] = torch.cat([edge_features['0'], r[..., None]], dim=1)
+    else:
+        edge_features['0'] = r[..., None]
+
+    return edge_features
+
+
+class SE3Transformer(nn.Module):
+    def __init__(self,
+                 num_layers: int,
+                 fiber_in: Fiber,
+                 fiber_hidden: Fiber,
+                 fiber_out: Fiber,
+                 num_heads: int,
+                 channels_div: int,
+                 fiber_edge: Fiber = Fiber({}),
+                 return_type: Optional[int] = None,
+                 pooling: Optional[Literal['avg', 'max']] = None,
+                 norm: bool = True,
+                 use_layer_norm: bool = True,
+                 tensor_cores: bool = False,
+                 low_memory: bool = False,
+                 **kwargs):
+        """
+        :param num_layers:          Number of attention layers
+        :param fiber_in:            Input fiber description
+        :param fiber_hidden:        Hidden fiber description
+        :param fiber_out:           Output fiber description
+        :param fiber_edge:          Input edge fiber description
+        :param num_heads:           Number of attention heads
+        :param channels_div:        Channels division before feeding to attention layer
+        :param return_type:         Return only features of this type
+        :param pooling:             'avg' or 'max' graph pooling before MLP layers
+        :param norm:                Apply a normalization layer after each attention block
+        :param use_layer_norm:      Apply layer normalization between MLP layers
+        :param tensor_cores:        True if using Tensor Cores (affects the use of fully fused convs, and padded bases)
+        :param low_memory:          If True, will use slower ops that use less memory
+        """
+        super().__init__()
+        self.num_layers = num_layers
+        self.fiber_edge = fiber_edge
+        self.num_heads = num_heads
+        self.channels_div = channels_div
+        self.return_type = return_type
+        self.pooling = pooling
+        self.max_degree = max(*fiber_in.degrees, *fiber_hidden.degrees, *fiber_out.degrees)
+        self.tensor_cores = tensor_cores
+        self.low_memory = low_memory
+
+        if low_memory and not tensor_cores:
+            logging.warning('Low memory mode will have no effect with no Tensor Cores')
+
+        # Fully fused convolutions when using Tensor Cores (and not low memory mode)
+        fuse_level = ConvSE3FuseLevel.FULL if tensor_cores and not low_memory else ConvSE3FuseLevel.PARTIAL
+
+        graph_modules = []
+        for i in range(num_layers):
+            graph_modules.append(AttentionBlockSE3(fiber_in=fiber_in,
+                                                   fiber_out=fiber_hidden,
+                                                   fiber_edge=fiber_edge,
+                                                   num_heads=num_heads,
+                                                   channels_div=channels_div,
+                                                   use_layer_norm=use_layer_norm,
+                                                   max_degree=self.max_degree,
+                                                   fuse_level=fuse_level))
+            if norm:
+                graph_modules.append(NormSE3(fiber_hidden))
+            fiber_in = fiber_hidden
+
+        graph_modules.append(ConvSE3(fiber_in=fiber_in,
+                                     fiber_out=fiber_out,
+                                     fiber_edge=fiber_edge,
+                                     self_interaction=True,
+                                     use_layer_norm=use_layer_norm,
+                                     max_degree=self.max_degree))
+        self.graph_modules = Sequential(*graph_modules)
+
+        if pooling is not None:
+            assert return_type is not None, 'return_type must be specified when pooling'
+            self.pooling_module = GPooling(pool=pooling, feat_type=return_type)
+
+    def forward(self, graph: DGLGraph, node_feats: Dict[str, Tensor],
+                edge_feats: Optional[Dict[str, Tensor]] = None,
+                basis: Optional[Dict[str, Tensor]] = None):
+        # Compute bases in case they weren't precomputed as part of the data loading
+        basis = basis or get_basis(graph.edata['rel_pos'], max_degree=self.max_degree, compute_gradients=False,
+                                   use_pad_trick=self.tensor_cores and not self.low_memory,
+                                   amp=torch.is_autocast_enabled())
+
+        # Add fused bases (per output degree, per input degree, and fully fused) to the dict
+        basis = update_basis_with_fused(basis, self.max_degree, use_pad_trick=self.tensor_cores and not self.low_memory,
+                                        fully_fused=self.tensor_cores and not self.low_memory)
+
+        edge_feats = get_populated_edge_features(graph.edata['rel_pos'], edge_feats)
+
+        node_feats = self.graph_modules(node_feats, edge_feats, graph=graph, basis=basis)
+
+        if self.pooling is not None:
+            return self.pooling_module(node_feats, graph=graph)
+
+        if self.return_type is not None:
+            return node_feats[str(self.return_type)]
+
+        return node_feats
+
+    @staticmethod
+    def add_argparse_args(parser):
+        parser.add_argument('--num_layers', type=int, default=7,
+                            help='Number of stacked Transformer layers')
+        parser.add_argument('--num_heads', type=int, default=8,
+                            help='Number of heads in self-attention')
+        parser.add_argument('--channels_div', type=int, default=2,
+                            help='Channels division before feeding to attention layer')
+        parser.add_argument('--pooling', type=str, default=None, const=None, nargs='?', choices=['max', 'avg'],
+                            help='Type of graph pooling')
+        parser.add_argument('--norm', type=str2bool, nargs='?', const=True, default=False,
+                            help='Apply a normalization layer after each attention block')
+        parser.add_argument('--use_layer_norm', type=str2bool, nargs='?', const=True, default=False,
+                            help='Apply layer normalization between MLP layers')
+        parser.add_argument('--low_memory', type=str2bool, nargs='?', const=True, default=False,
+                            help='If true, will use fused ops that are slower but that use less memory '
+                                 '(expect 25 percent less memory). '
+                                 'Only has an effect if AMP is enabled on Volta GPUs, or if running on Ampere GPUs')
+
+        return parser
+
+
+class SE3TransformerPooled(nn.Module):
+    def __init__(self,
+                 fiber_in: Fiber,
+                 fiber_out: Fiber,
+                 fiber_edge: Fiber,
+                 num_degrees: int,
+                 num_channels: int,
+                 output_dim: int,
+                 **kwargs):
+        super().__init__()
+        kwargs['pooling'] = kwargs['pooling'] or 'max'
+        self.transformer = SE3Transformer(
+            fiber_in=fiber_in,
+            fiber_hidden=Fiber.create(num_degrees, num_channels),
+            fiber_out=fiber_out,
+            fiber_edge=fiber_edge,
+            return_type=0,
+            **kwargs
+        )
+
+        n_out_features = fiber_out.num_features
+        self.mlp = nn.Sequential(
+            nn.Linear(n_out_features, n_out_features),
+            nn.ReLU(),
+            nn.Linear(n_out_features, output_dim)
+        )
+
+    def forward(self, graph, node_feats, edge_feats, basis=None):
+        feats = self.transformer(graph, node_feats, edge_feats, basis).squeeze(-1)
+        y = self.mlp(feats).squeeze(-1)
+        return y
+
+    @staticmethod
+    def add_argparse_args(parent_parser):
+        parser = parent_parser.add_argument_group("Model architecture")
+        SE3Transformer.add_argparse_args(parser)
+        parser.add_argument('--num_degrees',
+                            help='Number of degrees to use. Hidden features will have types [0, ..., num_degrees - 1]',
+                            type=int, default=4)
+        parser.add_argument('--num_channels', help='Number of channels for the hidden features', type=int, default=32)
+        return parent_parser

env/SE3Transformer/build/lib/se3_transformer/runtime/arguments.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import argparse
+import pathlib
+
+from se3_transformer.data_loading import QM9DataModule
+from se3_transformer.model import SE3TransformerPooled
+from se3_transformer.runtime.utils import str2bool
+
+PARSER = argparse.ArgumentParser(description='SE(3)-Transformer')
+
+paths = PARSER.add_argument_group('Paths')
+paths.add_argument('--data_dir', type=pathlib.Path, default=pathlib.Path('./data'),
+                   help='Directory where the data is located or should be downloaded')
+paths.add_argument('--log_dir', type=pathlib.Path, default=pathlib.Path('/results'),
+                   help='Directory where the results logs should be saved')
+paths.add_argument('--dllogger_name', type=str, default='dllogger_results.json',
+                   help='Name for the resulting DLLogger JSON file')
+paths.add_argument('--save_ckpt_path', type=pathlib.Path, default=None,
+                   help='File where the checkpoint should be saved')
+paths.add_argument('--load_ckpt_path', type=pathlib.Path, default=None,
+                   help='File of the checkpoint to be loaded')
+
+optimizer = PARSER.add_argument_group('Optimizer')
+optimizer.add_argument('--optimizer', choices=['adam', 'sgd', 'lamb'], default='adam')
+optimizer.add_argument('--learning_rate', '--lr', dest='learning_rate', type=float, default=0.002)
+optimizer.add_argument('--min_learning_rate', '--min_lr', dest='min_learning_rate', type=float, default=None)
+optimizer.add_argument('--momentum', type=float, default=0.9)
+optimizer.add_argument('--weight_decay', type=float, default=0.1)
+
+PARSER.add_argument('--epochs', type=int, default=100, help='Number of training epochs')
+PARSER.add_argument('--batch_size', type=int, default=240, help='Batch size')
+PARSER.add_argument('--seed', type=int, default=None, help='Set a seed globally')
+PARSER.add_argument('--num_workers', type=int, default=8, help='Number of dataloading workers')
+
+PARSER.add_argument('--amp', type=str2bool, nargs='?', const=True, default=False, help='Use Automatic Mixed Precision')
+PARSER.add_argument('--gradient_clip', type=float, default=None, help='Clipping of the gradient norms')
+PARSER.add_argument('--accumulate_grad_batches', type=int, default=1, help='Gradient accumulation')
+PARSER.add_argument('--ckpt_interval', type=int, default=-1, help='Save a checkpoint every N epochs')
+PARSER.add_argument('--eval_interval', dest='eval_interval', type=int, default=1,
+                    help='Do an evaluation round every N epochs')
+PARSER.add_argument('--silent', type=str2bool, nargs='?', const=True, default=False,
+                    help='Minimize stdout output')
+
+PARSER.add_argument('--benchmark', type=str2bool, nargs='?', const=True, default=False,
+                    help='Benchmark mode')
+
+QM9DataModule.add_argparse_args(PARSER)
+SE3TransformerPooled.add_argparse_args(PARSER)

env/SE3Transformer/build/lib/se3_transformer/runtime/callbacks.py

@@ -0,0 +1,160 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import logging
+import time
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import numpy as np
+import torch
+
+from se3_transformer.runtime.loggers import Logger
+from se3_transformer.runtime.metrics import MeanAbsoluteError
+
+
+class BaseCallback(ABC):
+    def on_fit_start(self, optimizer, args):
+        pass
+
+    def on_fit_end(self):
+        pass
+
+    def on_epoch_end(self):
+        pass
+
+    def on_batch_start(self):
+        pass
+
+    def on_validation_step(self, input, target, pred):
+        pass
+
+    def on_validation_end(self, epoch=None):
+        pass
+
+    def on_checkpoint_load(self, checkpoint):
+        pass
+
+    def on_checkpoint_save(self, checkpoint):
+        pass
+
+
+class LRSchedulerCallback(BaseCallback):
+    def __init__(self, logger: Optional[Logger] = None):
+        self.logger = logger
+        self.scheduler = None
+
+    @abstractmethod
+    def get_scheduler(self, optimizer, args):
+        pass
+
+    def on_fit_start(self, optimizer, args):
+        self.scheduler = self.get_scheduler(optimizer, args)
+
+    def on_checkpoint_load(self, checkpoint):
+        self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+
+    def on_checkpoint_save(self, checkpoint):
+        checkpoint['scheduler_state_dict'] = self.scheduler.state_dict()
+
+    def on_epoch_end(self):
+        if self.logger is not None:
+            self.logger.log_metrics({'learning rate': self.scheduler.get_last_lr()[0]}, step=self.scheduler.last_epoch)
+        self.scheduler.step()
+
+
+class QM9MetricCallback(BaseCallback):
+    """ Logs the rescaled mean absolute error for QM9 regression tasks """
+
+    def __init__(self, logger, targets_std, prefix=''):
+        self.mae = MeanAbsoluteError()
+        self.logger = logger
+        self.targets_std = targets_std
+        self.prefix = prefix
+        self.best_mae = float('inf')
+
+    def on_validation_step(self, input, target, pred):
+        self.mae(pred.detach(), target.detach())
+
+    def on_validation_end(self, epoch=None):
+        mae = self.mae.compute() * self.targets_std
+        logging.info(f'{self.prefix} MAE: {mae}')
+        self.logger.log_metrics({f'{self.prefix} MAE': mae}, epoch)
+        self.best_mae = min(self.best_mae, mae)
+
+    def on_fit_end(self):
+        if self.best_mae != float('inf'):
+            self.logger.log_metrics({f'{self.prefix} best MAE': self.best_mae})
+
+
+class QM9LRSchedulerCallback(LRSchedulerCallback):
+    def __init__(self, logger, epochs):
+        super().__init__(logger)
+        self.epochs = epochs
+
+    def get_scheduler(self, optimizer, args):
+        min_lr = args.min_learning_rate if args.min_learning_rate else args.learning_rate / 10.0
+        return torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, self.epochs, eta_min=min_lr)
+
+
+class PerformanceCallback(BaseCallback):
+    def __init__(self, logger, batch_size: int, warmup_epochs: int = 1, mode: str = 'train'):
+        self.batch_size = batch_size
+        self.warmup_epochs = warmup_epochs
+        self.epoch = 0
+        self.timestamps = []
+        self.mode = mode
+        self.logger = logger
+
+    def on_batch_start(self):
+        if self.epoch >= self.warmup_epochs:
+            self.timestamps.append(time.time() * 1000.0)
+
+    def _log_perf(self):
+        stats = self.process_performance_stats()
+        for k, v in stats.items():
+            logging.info(f'performance {k}: {v}')
+
+        self.logger.log_metrics(stats)
+
+    def on_epoch_end(self):
+        self.epoch += 1
+
+    def on_fit_end(self):
+        if self.epoch > self.warmup_epochs:
+            self._log_perf()
+            self.timestamps = []
+
+    def process_performance_stats(self):
+        timestamps = np.asarray(self.timestamps)
+        deltas = np.diff(timestamps)
+        throughput = (self.batch_size / deltas).mean()
+        stats = {
+            f"throughput_{self.mode}": throughput,
+            f"latency_{self.mode}_mean": deltas.mean(),
+            f"total_time_{self.mode}": timestamps[-1] - timestamps[0],
+        }
+        for level in [90, 95, 99]:
+            stats.update({f"latency_{self.mode}_{level}": np.percentile(deltas, level)})
+
+        return stats

env/SE3Transformer/build/lib/se3_transformer/runtime/gpu_affinity.py

@@ -0,0 +1,325 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import collections
+import itertools
+import math
+import os
+import pathlib
+import re
+
+import pynvml
+
+
+class Device:
+    # assumes nvml returns list of 64 bit ints
+    _nvml_affinity_elements = math.ceil(os.cpu_count() / 64)
+
+    def __init__(self, device_idx):
+        super().__init__()
+        self.handle = pynvml.nvmlDeviceGetHandleByIndex(device_idx)
+
+    def get_name(self):
+        return pynvml.nvmlDeviceGetName(self.handle)
+
+    def get_uuid(self):
+        return pynvml.nvmlDeviceGetUUID(self.handle)
+
+    def get_cpu_affinity(self):
+        affinity_string = ""
+        for j in pynvml.nvmlDeviceGetCpuAffinity(self.handle, Device._nvml_affinity_elements):
+            # assume nvml returns list of 64 bit ints
+            affinity_string = "{:064b}".format(j) + affinity_string
+
+        affinity_list = [int(x) for x in affinity_string]
+        affinity_list.reverse()  # so core 0 is in 0th element of list
+
+        ret = [i for i, e in enumerate(affinity_list) if e != 0]
+        return ret
+
+
+def get_thread_siblings_list():
+    """
+    Returns a list of 2-element integer tuples representing pairs of
+    hyperthreading cores.
+    """
+    path = "/sys/devices/system/cpu/cpu*/topology/thread_siblings_list"
+    thread_siblings_list = []
+    pattern = re.compile(r"(\d+)\D(\d+)")
+    for fname in pathlib.Path(path[0]).glob(path[1:]):
+        with open(fname) as f:
+            content = f.read().strip()
+            res = pattern.findall(content)
+            if res:
+                pair = tuple(map(int, res[0]))
+                thread_siblings_list.append(pair)
+    return thread_siblings_list
+
+
+def check_socket_affinities(socket_affinities):
+    # sets of cores should be either identical or disjoint
+    for i, j in itertools.product(socket_affinities, socket_affinities):
+        if not set(i) == set(j) and not set(i).isdisjoint(set(j)):
+            raise RuntimeError(f"Sets of cores should be either identical or disjoint, " f"but got {i} and {j}.")
+
+
+def get_socket_affinities(nproc_per_node, exclude_unavailable_cores=True):
+    devices = [Device(i) for i in range(nproc_per_node)]
+    socket_affinities = [dev.get_cpu_affinity() for dev in devices]
+
+    if exclude_unavailable_cores:
+        available_cores = os.sched_getaffinity(0)
+        socket_affinities = [list(set(affinity) & available_cores) for affinity in socket_affinities]
+
+    check_socket_affinities(socket_affinities)
+
+    return socket_affinities
+
+
+def set_socket_affinity(gpu_id):
+    """
+    The process is assigned with all available logical CPU cores from the CPU
+    socket connected to the GPU with a given id.
+
+    Args:
+        gpu_id: index of a GPU
+    """
+    dev = Device(gpu_id)
+    affinity = dev.get_cpu_affinity()
+    os.sched_setaffinity(0, affinity)
+
+
+def set_single_affinity(gpu_id):
+    """
+    The process is assigned with the first available logical CPU core from the
+    list of all CPU cores from the CPU socket connected to the GPU with a given
+    id.
+
+    Args:
+        gpu_id: index of a GPU
+    """
+    dev = Device(gpu_id)
+    affinity = dev.get_cpu_affinity()
+
+    # exclude unavailable cores
+    available_cores = os.sched_getaffinity(0)
+    affinity = list(set(affinity) & available_cores)
+    os.sched_setaffinity(0, affinity[:1])
+
+
+def set_single_unique_affinity(gpu_id, nproc_per_node):
+    """
+    The process is assigned with a single unique available physical CPU core
+    from the list of all CPU cores from the CPU socket connected to the GPU with
+    a given id.
+
+    Args:
+        gpu_id: index of a GPU
+    """
+    socket_affinities = get_socket_affinities(nproc_per_node)
+
+    siblings_list = get_thread_siblings_list()
+    siblings_dict = dict(siblings_list)
+
+    # remove siblings
+    for idx, socket_affinity in enumerate(socket_affinities):
+        socket_affinities[idx] = list(set(socket_affinity) - set(siblings_dict.values()))
+
+    affinities = []
+    assigned = []
+
+    for socket_affinity in socket_affinities:
+        for core in socket_affinity:
+            if core not in assigned:
+                affinities.append([core])
+                assigned.append(core)
+                break
+    os.sched_setaffinity(0, affinities[gpu_id])
+
+
+def set_socket_unique_affinity(gpu_id, nproc_per_node, mode, balanced=True):
+    """
+    The process is assigned with an unique subset of available physical CPU
+    cores from the CPU socket connected to a GPU with a given id.
+    Assignment automatically includes hyperthreading siblings (if siblings are
+    available).
+
+    Args:
+        gpu_id: index of a GPU
+        nproc_per_node: total number of processes per node
+        mode: mode
+        balanced: assign an equal number of physical cores to each process
+    """
+    socket_affinities = get_socket_affinities(nproc_per_node)
+
+    siblings_list = get_thread_siblings_list()
+    siblings_dict = dict(siblings_list)
+
+    # remove hyperthreading siblings
+    for idx, socket_affinity in enumerate(socket_affinities):
+        socket_affinities[idx] = list(set(socket_affinity) - set(siblings_dict.values()))
+
+    socket_affinities_to_device_ids = collections.defaultdict(list)
+
+    for idx, socket_affinity in enumerate(socket_affinities):
+        socket_affinities_to_device_ids[tuple(socket_affinity)].append(idx)
+
+    # compute minimal number of physical cores per GPU across all GPUs and
+    # sockets, code assigns this number of cores per GPU if balanced == True
+    min_physical_cores_per_gpu = min(
+        [len(cores) // len(gpus) for cores, gpus in socket_affinities_to_device_ids.items()]
+    )
+
+    for socket_affinity, device_ids in socket_affinities_to_device_ids.items():
+        devices_per_group = len(device_ids)
+        if balanced:
+            cores_per_device = min_physical_cores_per_gpu
+            socket_affinity = socket_affinity[: devices_per_group * min_physical_cores_per_gpu]
+        else:
+            cores_per_device = len(socket_affinity) // devices_per_group
+
+        for group_id, device_id in enumerate(device_ids):
+            if device_id == gpu_id:
+
+                # In theory there should be no difference in performance between
+                # 'interleaved' and 'continuous' pattern on Intel-based DGX-1,
+                # but 'continuous' should be better for DGX A100 because on AMD
+                # Rome 4 consecutive cores are sharing L3 cache.
+                # TODO: code doesn't attempt to automatically detect layout of
+                # L3 cache, also external environment may already exclude some
+                # cores, this code makes no attempt to detect it and to align
+                # mapping to multiples of 4.
+
+                if mode == "interleaved":
+                    affinity = list(socket_affinity[group_id::devices_per_group])
+                elif mode == "continuous":
+                    affinity = list(socket_affinity[group_id * cores_per_device: (group_id + 1) * cores_per_device])
+                else:
+                    raise RuntimeError("Unknown set_socket_unique_affinity mode")
+
+                # unconditionally reintroduce hyperthreading siblings, this step
+                # may result in a different numbers of logical cores assigned to
+                # each GPU even if balanced == True (if hyperthreading siblings
+                # aren't available for a subset of cores due to some external
+                # constraints, siblings are re-added unconditionally, in the
+                # worst case unavailable logical core will be ignored by
+                # os.sched_setaffinity().
+                affinity += [siblings_dict[aff] for aff in affinity if aff in siblings_dict]
+                os.sched_setaffinity(0, affinity)
+
+
+def set_affinity(gpu_id, nproc_per_node, mode="socket_unique_continuous", balanced=True):
+    """
+    The process is assigned with a proper CPU affinity which matches hardware
+    architecture on a given platform. Usually it improves and stabilizes
+    performance of deep learning training workloads.
+
+    This function assumes that the workload is running in multi-process
+    single-device mode (there are multiple training processes and each process
+    is running on a single GPU), which is typical for multi-GPU training
+    workloads using `torch.nn.parallel.DistributedDataParallel`.
+
+    Available affinity modes:
+    * 'socket' - the process is assigned with all available logical CPU cores
+    from the CPU socket connected to the GPU with a given id.
+    * 'single' - the process is assigned with the first available logical CPU
+    core from the list of all CPU cores from the CPU socket connected to the GPU
+    with a given id (multiple GPUs could be assigned with the same CPU core).
+    * 'single_unique' - the process is assigned with a single unique available
+    physical CPU core from the list of all CPU cores from the CPU socket
+    connected to the GPU with a given id.
+    * 'socket_unique_interleaved' - the process is assigned with an unique
+    subset of available physical CPU cores from the CPU socket connected to a
+    GPU with a given id, hyperthreading siblings are included automatically,
+    cores are assigned with interleaved indexing pattern
+    * 'socket_unique_continuous' - (the default) the process is assigned with an
+    unique subset of available physical CPU cores from the CPU socket connected
+    to a GPU with a given id, hyperthreading siblings are included
+    automatically, cores are assigned with continuous indexing pattern
+
+    'socket_unique_continuous' is the recommended mode for deep learning
+    training workloads on NVIDIA DGX machines.
+
+    Args:
+        gpu_id: integer index of a GPU
+        nproc_per_node: number of processes per node
+        mode: affinity mode
+        balanced: assign an equal number of physical cores to each process,
+            affects only 'socket_unique_interleaved' and
+            'socket_unique_continuous' affinity modes
+
+    Returns a set of logical CPU cores on which the process is eligible to run.
+
+    Example:
+
+    import argparse
+    import os
+
+    import gpu_affinity
+    import torch
+
+
+    def main():
+        parser = argparse.ArgumentParser()
+        parser.add_argument(
+            '--local_rank',
+            type=int,
+            default=os.getenv('LOCAL_RANK', 0),
+        )
+        args = parser.parse_args()
+
+        nproc_per_node = torch.cuda.device_count()
+
+        affinity = gpu_affinity.set_affinity(args.local_rank, nproc_per_node)
+        print(f'{args.local_rank}: core affinity: {affinity}')
+
+
+    if __name__ == "__main__":
+        main()
+
+    Launch the example with:
+    python -m torch.distributed.launch --nproc_per_node <#GPUs> example.py
+
+
+    WARNING: On DGX A100 only a half of CPU cores have direct access to GPUs.
+    This function restricts execution only to the CPU cores directly connected
+    to GPUs, so on DGX A100 it will limit the code to half of CPU cores and half
+    of CPU memory bandwidth (which may be fine for many DL models).
+    """
+    pynvml.nvmlInit()
+
+    if mode == "socket":
+        set_socket_affinity(gpu_id)
+    elif mode == "single":
+        set_single_affinity(gpu_id)
+    elif mode == "single_unique":
+        set_single_unique_affinity(gpu_id, nproc_per_node)
+    elif mode == "socket_unique_interleaved":
+        set_socket_unique_affinity(gpu_id, nproc_per_node, "interleaved", balanced)
+    elif mode == "socket_unique_continuous":
+        set_socket_unique_affinity(gpu_id, nproc_per_node, "continuous", balanced)
+    else:
+        raise RuntimeError("Unknown affinity mode")
+
+    affinity = os.sched_getaffinity(0)
+    return affinity

env/SE3Transformer/build/lib/se3_transformer/runtime/inference.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+from typing import List
+
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from se3_transformer.runtime import gpu_affinity
+from se3_transformer.runtime.arguments import PARSER
+from se3_transformer.runtime.callbacks import BaseCallback
+from se3_transformer.runtime.loggers import DLLogger
+from se3_transformer.runtime.utils import to_cuda, get_local_rank
+
+
+@torch.inference_mode()
+def evaluate(model: nn.Module,
+             dataloader: DataLoader,
+             callbacks: List[BaseCallback],
+             args):
+    model.eval()
+    for i, batch in tqdm(enumerate(dataloader), total=len(dataloader), unit='batch', desc=f'Evaluation',
+                         leave=False, disable=(args.silent or get_local_rank() != 0)):
+        *input, target = to_cuda(batch)
+
+        for callback in callbacks:
+            callback.on_batch_start()
+
+        with torch.cuda.amp.autocast(enabled=args.amp):
+            pred = model(*input)
+
+            for callback in callbacks:
+                callback.on_validation_step(input, target, pred)
+
+
+if __name__ == '__main__':
+    from se3_transformer.runtime.callbacks import QM9MetricCallback, PerformanceCallback
+    from se3_transformer.runtime.utils import init_distributed, seed_everything
+    from se3_transformer.model import SE3TransformerPooled, Fiber
+    from se3_transformer.data_loading import QM9DataModule
+    import torch.distributed as dist
+    import logging
+    import sys
+
+    is_distributed = init_distributed()
+    local_rank = get_local_rank()
+    args = PARSER.parse_args()
+
+    logging.getLogger().setLevel(logging.CRITICAL if local_rank != 0 or args.silent else logging.INFO)
+
+    logging.info('====== SE(3)-Transformer ======')
+    logging.info('|  Inference on the test set  |')
+    logging.info('===============================')
+
+    if not args.benchmark and args.load_ckpt_path is None:
+        logging.error('No load_ckpt_path provided, you need to provide a saved model to evaluate')
+        sys.exit(1)
+
+    if args.benchmark:
+        logging.info('Running benchmark mode with one warmup pass')
+
+    if args.seed is not None:
+        seed_everything(args.seed)
+
+    major_cc, minor_cc = torch.cuda.get_device_capability()
+
+    logger = DLLogger(args.log_dir, filename=args.dllogger_name)
+    datamodule = QM9DataModule(**vars(args))
+    model = SE3TransformerPooled(
+        fiber_in=Fiber({0: datamodule.NODE_FEATURE_DIM}),
+        fiber_out=Fiber({0: args.num_degrees * args.num_channels}),
+        fiber_edge=Fiber({0: datamodule.EDGE_FEATURE_DIM}),
+        output_dim=1,
+        tensor_cores=(args.amp and major_cc >= 7) or major_cc >= 8,  # use Tensor Cores more effectively
+        **vars(args)
+    )
+    callbacks = [QM9MetricCallback(logger, targets_std=datamodule.targets_std, prefix='test')]
+
+    model.to(device=torch.cuda.current_device())
+    if args.load_ckpt_path is not None:
+        checkpoint = torch.load(str(args.load_ckpt_path), map_location={'cuda:0': f'cuda:{local_rank}'})
+        model.load_state_dict(checkpoint['state_dict'])
+
+    if is_distributed:
+        nproc_per_node = torch.cuda.device_count()
+        affinity = gpu_affinity.set_affinity(local_rank, nproc_per_node)
+        model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
+
+    test_dataloader = datamodule.test_dataloader() if not args.benchmark else datamodule.train_dataloader()
+    evaluate(model,
+             test_dataloader,
+             callbacks,
+             args)
+
+    for callback in callbacks:
+        callback.on_validation_end()
+
+    if args.benchmark:
+        world_size = dist.get_world_size() if dist.is_initialized() else 1
+        callbacks = [PerformanceCallback(logger, args.batch_size * world_size, warmup_epochs=1, mode='inference')]
+        for _ in range(6):
+            evaluate(model,
+                     test_dataloader,
+                     callbacks,
+                     args)
+            callbacks[0].on_epoch_end()
+
+        callbacks[0].on_fit_end()

env/SE3Transformer/build/lib/se3_transformer/runtime/loggers.py

@@ -0,0 +1,134 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+import pathlib
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Dict, Any, Callable, Optional
+
+import dllogger
+import torch.distributed as dist
+import wandb
+from dllogger import Verbosity
+
+from se3_transformer.runtime.utils import rank_zero_only
+
+
+class Logger(ABC):
+    @rank_zero_only
+    @abstractmethod
+    def log_hyperparams(self, params):
+        pass
+
+    @rank_zero_only
+    @abstractmethod
+    def log_metrics(self, metrics, step=None):
+        pass
+
+    @staticmethod
+    def _sanitize_params(params):
+        def _sanitize(val):
+            if isinstance(val, Callable):
+                try:
+                    _val = val()
+                    if isinstance(_val, Callable):
+                        return val.__name__
+                    return _val
+                except Exception:
+                    return getattr(val, "__name__", None)
+            elif isinstance(val, pathlib.Path) or isinstance(val, Enum):
+                return str(val)
+            return val
+
+        return {key: _sanitize(val) for key, val in params.items()}
+
+
+class LoggerCollection(Logger):
+    def __init__(self, loggers):
+        super().__init__()
+        self.loggers = loggers
+
+    def __getitem__(self, index):
+        return [logger for logger in self.loggers][index]
+
+    @rank_zero_only
+    def log_metrics(self, metrics, step=None):
+        for logger in self.loggers:
+            logger.log_metrics(metrics, step)
+
+    @rank_zero_only
+    def log_hyperparams(self, params):
+        for logger in self.loggers:
+            logger.log_hyperparams(params)
+
+
+class DLLogger(Logger):
+    def __init__(self, save_dir: pathlib.Path, filename: str):
+        super().__init__()
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            save_dir.mkdir(parents=True, exist_ok=True)
+            dllogger.init(
+                backends=[dllogger.JSONStreamBackend(Verbosity.DEFAULT, str(save_dir / filename))])
+
+    @rank_zero_only
+    def log_hyperparams(self, params):
+        params = self._sanitize_params(params)
+        dllogger.log(step="PARAMETER", data=params)
+
+    @rank_zero_only
+    def log_metrics(self, metrics, step=None):
+        if step is None:
+            step = tuple()
+
+        dllogger.log(step=step, data=metrics)
+
+
+class WandbLogger(Logger):
+    def __init__(
+            self,
+            name: str,
+            save_dir: pathlib.Path,
+            id: Optional[str] = None,
+            project: Optional[str] = None
+    ):
+        super().__init__()
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            save_dir.mkdir(parents=True, exist_ok=True)
+            self.experiment = wandb.init(name=name,
+                                         project=project,
+                                         id=id,
+                                         dir=str(save_dir),
+                                         resume='allow',
+                                         anonymous='must')
+
+    @rank_zero_only
+    def log_hyperparams(self, params: Dict[str, Any]) -> None:
+        params = self._sanitize_params(params)
+        self.experiment.config.update(params, allow_val_change=True)
+
+    @rank_zero_only
+    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
+        if step is not None:
+            self.experiment.log({**metrics, 'epoch': step})
+        else:
+            self.experiment.log(metrics)

env/SE3Transformer/build/lib/se3_transformer/runtime/metrics.py

@@ -0,0 +1,83 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+#
+# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES
+# SPDX-License-Identifier: MIT
+
+from abc import ABC, abstractmethod
+
+import torch
+import torch.distributed as dist
+from torch import Tensor
+
+
+class Metric(ABC):
+    """ Metric class with synchronization capabilities similar to TorchMetrics """
+
+    def __init__(self):
+        self.states = {}
+
+    def add_state(self, name: str, default: Tensor):
+        assert name not in self.states
+        self.states[name] = default.clone()
+        setattr(self, name, default)
+
+    def synchronize(self):
+        if dist.is_initialized():
+            for state in self.states:
+                dist.all_reduce(getattr(self, state), op=dist.ReduceOp.SUM, group=dist.group.WORLD)
+
+    def __call__(self, *args, **kwargs):
+        self.update(*args, **kwargs)
+
+    def reset(self):
+        for name, default in self.states.items():
+            setattr(self, name, default.clone())
+
+    def compute(self):
+        self.synchronize()
+        value = self._compute().item()
+        self.reset()
+        return value
+
+    @abstractmethod
+    def _compute(self):
+        pass
+
+    @abstractmethod
+    def update(self, preds: Tensor, targets: Tensor):
+        pass
+
+
+class MeanAbsoluteError(Metric):
+    def __init__(self):
+        super().__init__()
+        self.add_state('error', torch.tensor(0, dtype=torch.float32, device='cuda'))
+        self.add_state('total', torch.tensor(0, dtype=torch.int32, device='cuda'))
+
+    def update(self, preds: Tensor, targets: Tensor):
+        preds = preds.detach()
+        n = preds.shape[0]
+        error = torch.abs(preds.view(n, -1) - targets.view(n, -1)).sum()
+        self.total += n
+        self.error += error
+
+    def _compute(self):
+        return self.error / self.total