Upload 187 files

.gitattributes

@@ -33,3 +33,17 @@ 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
+assets/overview.png filter=lfs diff=lfs merge=lfs -text
+ShapeID/out/2d/image_with_v.png filter=lfs diff=lfs merge=lfs -text
+ShapeID/out/2d/progression/New[[:space:]]Folder[[:space:]]With[[:space:]]Items/0.png filter=lfs diff=lfs merge=lfs -text
+ShapeID/out/2d/progression/New[[:space:]]Folder[[:space:]]With[[:space:]]Items/1.png filter=lfs diff=lfs merge=lfs -text
+ShapeID/out/2d/progression/New[[:space:]]Folder[[:space:]]With[[:space:]]Items/10.png filter=lfs diff=lfs merge=lfs -text
+ShapeID/out/2d/progression/New[[:space:]]Folder[[:space:]]With[[:space:]]Items/2.png filter=lfs diff=lfs merge=lfs -text
+ShapeID/out/2d/progression/New[[:space:]]Folder[[:space:]]With[[:space:]]Items/3.png filter=lfs diff=lfs merge=lfs -text
+ShapeID/out/2d/progression/New[[:space:]]Folder[[:space:]]With[[:space:]]Items/4.png filter=lfs diff=lfs merge=lfs -text
+ShapeID/out/2d/progression/New[[:space:]]Folder[[:space:]]With[[:space:]]Items/5.png filter=lfs diff=lfs merge=lfs -text
+ShapeID/out/2d/progression/New[[:space:]]Folder[[:space:]]With[[:space:]]Items/6.png filter=lfs diff=lfs merge=lfs -text
+ShapeID/out/2d/progression/New[[:space:]]Folder[[:space:]]With[[:space:]]Items/7.png filter=lfs diff=lfs merge=lfs -text
+ShapeID/out/2d/progression/New[[:space:]]Folder[[:space:]]With[[:space:]]Items/8.png filter=lfs diff=lfs merge=lfs -text
+ShapeID/out/2d/progression/New[[:space:]]Folder[[:space:]]With[[:space:]]Items/9.png filter=lfs diff=lfs merge=lfs -text
+ShapeID/out/2d/V.png filter=lfs diff=lfs merge=lfs -text

Generator/__init__.py

@@ -0,0 +1,22 @@
+
+"""
+Datasets interface.
+"""
+from .constants import dataset_setups
+from .datasets import BaseGen, BrainIDGen
+
+
+
+dataset_options = {  
+    'default': BaseGen,
+    'brain_id': BrainIDGen,
+}
+
+
+
+
+def build_datasets(gen_args, device):
+    """Helper function to build dataset for different splits ('train' or 'test')."""
+    datasets = {'all': dataset_options[gen_args.dataset_option](gen_args, device)}
+    return datasets
+

Generator/config.py

@@ -0,0 +1,181 @@
+
+"""Config utilities for yml file."""
+import os
+from argparse import Namespace
+import collections
+import functools
+import os
+import re
+
+import yaml
+# from imaginaire.utils.distributed import master_only_print as print
+
+
+class AttrDict(dict):
+    """Dict as attribute trick."""
+
+    def __init__(self, *args, **kwargs):
+        super(AttrDict, self).__init__(*args, **kwargs)
+        self.__dict__ = self
+        for key, value in self.__dict__.items():
+            if isinstance(value, dict):
+                self.__dict__[key] = AttrDict(value)
+            elif isinstance(value, (list, tuple)):
+                if isinstance(value[0], dict):
+                    self.__dict__[key] = [AttrDict(item) for item in value]
+                else:
+                    self.__dict__[key] = value
+
+    def yaml(self):
+        """Convert object to yaml dict and return."""
+        yaml_dict = {}
+        for key, value in self.__dict__.items():
+            if isinstance(value, AttrDict):
+                yaml_dict[key] = value.yaml()
+            elif isinstance(value, list):
+                if isinstance(value[0], AttrDict):
+                    new_l = []
+                    for item in value:
+                        new_l.append(item.yaml())
+                    yaml_dict[key] = new_l
+                else:
+                    yaml_dict[key] = value
+            else:
+                yaml_dict[key] = value
+        return yaml_dict
+
+    def __repr__(self):
+        """Print all variables."""
+        ret_str = []
+        for key, value in self.__dict__.items():
+            if isinstance(value, AttrDict):
+                ret_str.append('{}:'.format(key))
+                child_ret_str = value.__repr__().split('\n')
+                for item in child_ret_str:
+                    ret_str.append('    ' + item)
+            elif isinstance(value, list):
+                if isinstance(value[0], AttrDict):
+                    ret_str.append('{}:'.format(key))
+                    for item in value:
+                        # Treat as AttrDict above.
+                        child_ret_str = item.__repr__().split('\n')
+                        for item in child_ret_str:
+                            ret_str.append('    ' + item)
+                else:
+                    ret_str.append('{}: {}'.format(key, value))
+            else:
+                ret_str.append('{}: {}'.format(key, value))
+        return '\n'.join(ret_str)
+
+
+class Config(AttrDict):
+    r"""Configuration class. This should include every human specifiable
+    hyperparameter values for your training."""
+
+    def __init__(self, filename=None, verbose=False):
+        super(Config, self).__init__()
+
+        # Update with given configurations.
+        if os.path.exists(filename):
+
+            loader = yaml.SafeLoader
+            loader.add_implicit_resolver(
+                u'tag:yaml.org,2002:float',
+                re.compile(u'''^(?:
+                [-+]?(?:[0-9][0-9_]*)\\.[0-9_]*(?:[eE][-+]?[0-9]+)?
+                |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
+                |\\.[0-9_]+(?:[eE][-+][0-9]+)?
+                |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\.[0-9_]*
+                |[-+]?\\.(?:inf|Inf|INF)
+                |\\.(?:nan|NaN|NAN))$''', re.X),
+                list(u'-+0123456789.'))
+            try:
+                with open(filename, 'r') as f:
+                    cfg_dict = yaml.load(f, Loader=loader)
+            except EnvironmentError:
+                print('Please check the file with name of "%s"', filename)
+            recursive_update(self, cfg_dict)
+        else:
+            raise ValueError('Provided config path not existed: %s' % filename)
+
+        if verbose:
+            print(' imaginaire config '.center(80, '-'))
+            print(self.__repr__())
+            print(''.center(80, '-'))
+
+
+def rsetattr(obj, attr, val):
+    """Recursively find object and set value"""
+    pre, _, post = attr.rpartition('.')
+    return setattr(rgetattr(obj, pre) if pre else obj, post, val)
+
+
+def rgetattr(obj, attr, *args):
+    """Recursively find object and return value"""
+
+    def _getattr(obj, attr):
+        r"""Get attribute."""
+        return getattr(obj, attr, *args)
+
+    return functools.reduce(_getattr, [obj] + attr.split('.'))
+
+
+def recursive_update(d, u):
+    """Recursively update AttrDict d with AttrDict u"""
+    if u is not None:
+        for key, value in u.items():
+            if isinstance(value, collections.abc.Mapping):
+                d.__dict__[key] = recursive_update(d.get(key, AttrDict({})), value)
+            elif isinstance(value, (list, tuple)):
+                if len(value) > 0 and isinstance(value[0], dict):
+                    d.__dict__[key] = [AttrDict(item) for item in value]
+                else:
+                    d.__dict__[key] = value
+            else:
+                d.__dict__[key] = value
+    return d
+
+
+def merge_and_update_from_dict(cfg, dct):
+    """
+    (Compatible for submitit's Dict as attribute trick)
+    Merge dict as dict() to config as CfgNode().
+    Args:
+        cfg: dict
+        dct: dict
+    """
+    if dct is not None:
+        for key, value in dct.items():
+            if isinstance(value, dict):
+                if key in cfg.keys():
+                    sub_cfgnode = cfg[key]
+                else:
+                    sub_cfgnode = dict()
+                    cfg.__setattr__(key, sub_cfgnode) 
+                sub_cfgnode = merge_and_update_from_dict(sub_cfgnode, value)
+            else:
+                cfg[key] = value
+    return cfg
+
+
+def load_config(cfg_files = [], cfg_dir = ''):
+    cfg = Config(cfg_files[0]) 
+    for cfg_file in cfg_files[1:]:
+        add_cfg = Config(cfg_file)
+        cfg = merge_and_update_from_dict(cfg, add_cfg)
+    return cfg
+    
+
+def nested_dict_to_namespace(dictionary):
+    namespace = dictionary
+    if isinstance(dictionary, dict):
+        namespace = Namespace(**dictionary)
+        for key, value in dictionary.items():
+            setattr(namespace, key, nested_dict_to_namespace(value))
+    return namespace
+
+
+def preprocess_cfg(cfg_files, cfg_dir = ''):
+    config = load_config(cfg_files, cfg_dir)
+    args = nested_dict_to_namespace(config)
+    return args

Generator/constants.py

@@ -0,0 +1,290 @@
+import os, glob
+
+from .utils import *
+
+augmentation_funcs = {
+    'gamma': add_gamma_transform,
+    'bias_field': add_bias_field,
+    'resample': resample_resolution,
+    'noise': add_noise,
+}
+
+processing_funcs = {
+    'T1': read_and_deform_image,
+    'T2': read_and_deform_image,
+    'FLAIR': read_and_deform_image,
+    'CT': read_and_deform_CT,
+    'segmentation': read_and_deform_segmentation,
+    'surface': read_and_deform_surface,
+    'distance': read_and_deform_distance,
+    'bias_field': read_and_deform_bias_field,
+    'registration': read_and_deform_registration,
+    'pathology': read_and_deform_pathology, 
+}
+
+
+dataset_setups = { 
+
+    'ADHD': { 
+        'root': '/autofs/space/yogurt_001/users/pl629/data/adhd200_crop',
+        'pathology_type': None,
+        'train': 'train.txt',
+        'test': 'test.txt',
+        'modalities': ['T1'],
+
+        'paths':{
+                # for synth
+                'Gen': 'label_maps_generation', 
+                'Dmaps': None, 
+                'DmapsBag': None, 
+
+                # real images
+                'T1': 'T1', 
+                'T2': None, 
+                'FLAIR': None,
+                'CT': None,
+
+                # processed ground truths 
+                'surface': None, #'surfaces',  TODO
+                'distance': None,  
+                'segmentation': 'label_maps_segmentation',
+                'bias_field': None,
+                'pathology': None,
+                'pathology_prob': None,
+        }
+    },
+
+    'HCP': { 
+        'root': '/autofs/space/yogurt_001/users/pl629/data/hcp_crop',
+        'pathology_type': None,
+        'train': 'train.txt',
+        'test': 'test.txt',
+        'modalities': ['T1', 'T2'],
+
+        'paths':{
+                # for synth
+                'Gen': 'label_maps_generation', 
+                'Dmaps': None, 
+                'DmapsBag': None, 
+
+                # real images
+                'T1': 'T1', 
+                'T2': 'T2', 
+                'FLAIR': None,
+                'CT': None,
+
+                # processed ground truths 
+                'surface': None, #'surfaces', 
+                'distance': None,  
+                'segmentation': 'label_maps_segmentation',
+                'bias_field': None,
+                'pathology': None,
+                'pathology_prob': None,
+        }
+    },
+
+    'AIBL': { 
+        'root': '/autofs/space/yogurt_001/users/pl629/data/aibl_crop',
+        'pathology_type': None,
+        'train': 'train.txt',
+        'test': 'test.txt',
+        'modalities': ['T1', 'T2', 'FLAIR'],
+
+        'paths':{
+                # for synth
+                'Gen': 'label_maps_generation', 
+                'Dmaps': None, 
+                'DmapsBag': None, 
+
+                # real images
+                'T1': 'T1', 
+                'T2': 'T2', 
+                'FLAIR': 'FLAIR',
+                'CT': None,
+
+                # processed ground truths 
+                'surface': None, #'surfaces', 
+                'distance': None,  
+                'segmentation': 'label_maps_segmentation',
+                'bias_field': None,
+                'pathology': None,
+                'pathology_prob': None,
+        }
+    },
+
+    'OASIS': { 
+        'root': '/autofs/space/yogurt_001/users/pl629/data/oasis3',
+        'pathology_type': None,
+        'train': 'train.txt',
+        'test': 'test.txt',
+        'modalities': ['T1', 'CT'],
+
+        'paths':{
+                # for synth
+                'Gen': 'label_maps_generation', 
+                'Dmaps': None, 
+                'DmapsBag': None, 
+
+                # real images
+                'T1': 'T1', 
+                'T2': None, 
+                'FLAIR': None,
+                'CT': 'CT',
+
+                # processed ground truths 
+                'surface': None, #'surfaces', 
+                'distance': None,  
+                'segmentation': 'label_maps_segmentation',
+                'bias_field': None,
+                'pathology': None,
+                'pathology_prob': None,
+        }
+    },
+
+    'ADNI': { 
+        'root': '/autofs/space/yogurt_001/users/pl629/data/adni_crop',
+        'pathology_type': None, #'wmh',
+        'train': 'train.txt',
+        'test': 'test.txt',
+        'modalities': ['T1'],
+
+        'paths':{
+                # for synth
+                'Gen': 'label_maps_generation', 
+                'Dmaps': 'Dmaps', 
+                'DmapsBag': 'DmapsBag', 
+
+                # real images
+                'T1': 'T1', 
+                'T2': None, 
+                'FLAIR': None,
+                'CT': None,
+
+                # processed ground truths
+                'surface': 'surfaces',  
+                'distance': 'Dmaps',  
+                'segmentation': 'label_maps_segmentation',
+                'bias_field': None,
+                'pathology': 'pathology_maps_segmentation',
+                'pathology_prob': 'pathology_probability',
+        }
+    },
+
+    'ADNI3': { 
+        'root': '/autofs/space/yogurt_001/users/pl629/data/adni3_crop',
+        'pathology_type': None, # 'wmh',
+        'train': 'train.txt',
+        'test': 'test.txt',
+        'modalities': ['T1', 'FLAIR'],
+
+        'paths':{
+                # for synth
+                'Gen': 'label_maps_generation', 
+                'Dmaps': None, 
+                'DmapsBag': None, 
+
+                # real images
+                'T1': 'T1', 
+                'T2': None, 
+                'FLAIR': 'FLAIR',
+                'CT': None,
+
+                # processed ground truths 
+                'surface': None, #'surfaces',  TODO
+                'distance': None,  
+                'segmentation': 'label_maps_segmentation',
+                'bias_field': None,
+                'pathology': 'pathology_maps_segmentation',
+                'pathology_prob': 'pathology_probability',
+        }
+    },
+
+    'ATLAS': { 
+        'root': '/autofs/space/yogurt_001/users/pl629/data/atlas_crop',
+        'pathology_type': 'stroke',
+        'train': 'train.txt',
+        'test': 'test.txt',
+        'modalities': ['T1'],
+
+        'paths':{
+                # for synth
+                'Gen': 'label_maps_generation', 
+                'Dmaps': None, 
+                'DmapsBag': None, 
+
+                # real images
+                'T1': 'T1', 
+                'T2': None, 
+                'FLAIR': None,
+                'CT': None,
+
+                # processed ground truths 
+                'surface': None, #'surfaces',  TODO
+                'distance': None,  
+                'segmentation': 'label_maps_segmentation',
+                'bias_field': None,
+                'pathology': 'pathology_maps_segmentation',
+                'pathology_prob': 'pathology_probability',
+        }
+    },
+
+    'ISLES': { 
+        'root': '/autofs/space/yogurt_001/users/pl629/data/isles2022_crop',
+        'pathology_type': 'stroke',
+        'train': 'train.txt',
+        'test': 'test.txt',
+        'modalities': ['FLAIR'],
+
+        'paths':{
+                # for synth
+                'Gen': 'label_maps_generation', 
+                'Dmaps': None, 
+                'DmapsBag': None, 
+
+                # real images
+                'T1': None, 
+                'T2': None, 
+                'FLAIR': 'FLAIR',
+                'CT': None,
+
+                # processed ground truths 
+                'surface': None, #'surfaces',  TODO
+                'distance': None,  
+                'segmentation': 'label_maps_segmentation',
+                'bias_field': None,
+                'pathology': 'pathology_maps_segmentation',
+                'pathology_prob': 'pathology_probability',
+        }
+    },
+}
+
+
+all_dataset_names = dataset_setups.keys()
+
+
+# get all pathologies
+pathology_paths = []
+pathology_prob_paths = []
+for name, dict in dataset_setups.items():
+    # TODO: select what kind of shapes?
+    if dict['paths']['pathology'] is not None and dict['pathology_type'] is not None and dict['pathology_type'] == 'stroke':   
+        pathology_paths += glob.glob(os.path.join(dict['root'], dict['paths']['pathology'], '*.nii.gz')) \
+                                        + glob.glob(os.path.join(dict['root'], dict['paths']['pathology'], '*.nii'))
+        pathology_prob_paths += glob.glob(os.path.join(dict['root'], dict['paths']['pathology_prob'], '*.nii.gz')) \
+                                        + glob.glob(os.path.join(dict['root'], dict['paths']['pathology_prob'], '*.nii'))
+n_pathology = len(pathology_paths)
+
+
+# with csf # NOTE old version (FreeSurfer standard), non-vast
+label_list_segmentation = [0,14,15,16,24,77,85,   2, 3, 4, 7, 8, 10,11,12,13,17,18,26,28,   41,42,43,46,47,49,50,51,52,53,54,58,60] # 33
+n_neutral_labels = 7
+
+
+## NEW VAST synth
+label_list_segmentation_brainseg_with_extracerebral =  [0, 11, 12, 13, 16, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46,
+                                   1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 14, 15, 17, 47, 49, 51, 53, 55,
+                                   18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 48, 50, 52, 54, 56] 
+n_neutral_labels_brainseg_with_extracerebral = 20
+
+label_list_segmentation_brainseg_left = [0, 1, 2, 3, 4, 7, 8, 9, 10, 14, 15, 17, 31, 34, 36, 38, 40, 42]
+

Generator/datasets.py

@@ -0,0 +1,757 @@
+import os, sys, glob
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from collections import defaultdict 
+import random
+
+import torch
+import numpy as np
+import nibabel as nib
+from torch.utils.data import Dataset 
+
+
+from .utils import *
+from .constants import n_pathology, pathology_paths, pathology_prob_paths, \
+    n_neutral_labels_brainseg_with_extracerebral, label_list_segmentation_brainseg_with_extracerebral, \
+    label_list_segmentation_brainseg_left, augmentation_funcs, processing_funcs
+import utils.interpol as interpol
+
+from utils.misc import viewVolume
+
+
+from ShapeID.DiffEqs.pde import AdvDiffPDE
+ 
+
+
+class BaseGen(Dataset):
+    """
+    BaseGen dataset
+    """ 
+    def __init__(self, gen_args, device='cpu'):
+
+        self.gen_args = gen_args 
+        self.split = gen_args.split 
+
+        self.synth_args = self.gen_args.generator
+        self.shape_gen_args = gen_args.pathology_shape_generator
+        self.real_image_args = gen_args.real_image_generator
+        self.synth_image_args = gen_args.synth_image_generator 
+        self.augmentation_steps = vars(gen_args.augmentation_steps)
+        self.input_prob = vars(gen_args.modality_probs)
+        self.device = device
+
+        self.prepare_tasks()
+        self.prepare_paths()
+        self.prepare_grid()
+        self.prepare_one_hot()
+
+
+    def __len__(self):
+        return sum([len(self.names[i]) for i in range(len(self.names))])
+
+
+    def idx_to_path(self, idx):
+        cnt = 0
+        for i, l in enumerate(self.datasets_len):
+            if idx >= cnt and idx < cnt + l:
+                dataset_name = self.datasets[i]
+                age = self.ages[i][os.path.basename(self.names[i][idx - cnt]).split('.T1w')[0]] if len(self.ages) > 0 else None
+                return dataset_name, vars(self.input_prob[dataset_name]), self.names[i][idx - cnt], age
+            else:
+                cnt += l
+
+
+    def prepare_paths(self):
+
+        # Collect list of available images, per dataset
+        if len(self.gen_args.dataset_names) < 1:
+            datasets = [] 
+            g = glob.glob(os.path.join(self.gen_args.data_root, '*' + 'T1w.nii'))
+            for i in range(len(g)):
+                filename = os.path.basename(g[i])
+                dataset = filename[:filename.find('.')]
+                found = False
+                for d in datasets:
+                    if dataset == d:
+                        found = True
+                if found is False:
+                    datasets.append(dataset)
+            print('Found ' + str(len(datasets)) + ' datasets with ' + str(len(g)) + ' scans in total')
+        else:
+            datasets = self.gen_args.dataset_names
+        print('Dataset list', datasets)
+
+
+        names = [] 
+        if 'age' in self.tasks: 
+            self.split = self.split + '_age'
+        if self.gen_args.split_root is not None:
+            split_file = open(os.path.join(self.gen_args.split_root, self.split + '.txt'), 'r')
+            split_names = []
+            for subj in split_file.readlines():
+                split_names.append(subj.strip())  
+
+            for i in range(len(datasets)):
+                names.append([name for name in split_names if os.path.basename(name).startswith(datasets[i])]) 
+        #else:
+        #    for i in range(len(datasets)):
+        #        names.append(glob.glob(os.path.join(self.gen_args.data_root, datasets[i] + '.*' + 'T1w.nii')))
+
+        # read brain age
+        ages = []
+        if 'age' in self.tasks: 
+            age_file = open(os.path.join(self.gen_args.split_root, 'participants_age.txt'), 'r') 
+            subj_name_age = [] 
+            for line in age_file.readlines(): # 'subj age\n' 
+                subj_name_age.append(line.strip().split(' '))
+            for i in range(len(datasets)):
+                ages.append({})
+                for [name, age] in subj_name_age:
+                    if name.startswith(datasets[i]):
+                        ages[-1][name] = float(age)
+            print('Age info', self.split, len(ages[0].items()), min(ages[0].values()), max(ages[0].values()))
+            
+        self.ages = ages
+        self.names = names
+        self.datasets = datasets
+        self.datasets_num = len(datasets)
+        self.datasets_len = [len(self.names[i]) for i in range(len(self.names))]
+        print('Num of data', sum([len(self.names[i]) for i in range(len(self.names))]))
+
+        self.pathology_type = None #setup_dict['pathology_type']
+        
+
+    def prepare_tasks(self):
+        self.tasks = [key for (key, value) in vars(self.gen_args.task).items() if value]
+        if 'bias_field' in self.tasks and 'segmentation' not in self.tasks:
+            # add segmentation mask for computing bias_field_soft_mask
+            self.tasks += ['segmentation']
+        if 'pathology' in self.tasks and self.synth_args.augment_pathology and self.synth_args.random_shape_prob < 1.: 
+            self.t = torch.from_numpy(np.arange(self.shape_gen_args.max_nt) * self.shape_gen_args.dt).to(self.device)
+            with torch.no_grad():
+                self.adv_pde = AdvDiffPDE(data_spacing=[1., 1., 1.], 
+                                        perf_pattern='adv', 
+                                        V_type='vector_div_free', 
+                                        V_dict={},
+                                        BC=self.shape_gen_args.bc, 
+                                        dt=self.shape_gen_args.dt, 
+                                        device=self.device
+                                        )
+        else:
+            self.t, self.adv_pde = None, None
+        for task_name in self.tasks: 
+            if task_name not in processing_funcs.keys(): 
+                print('Warning: Function for task "%s" not found' % task_name)
+
+
+    def prepare_grid(self): 
+        self.size = self.synth_args.size
+
+        # Get resolution of training data
+        #aff = nib.load(os.path.join(self.modalities['Gen'], self.names[0])).affine
+        #self.res_training_data = np.sqrt(np.sum(abs(aff[:-1, :-1]), axis=0))
+
+        self.res_training_data = np.array([1.0, 1.0, 1.0])
+
+        xx, yy, zz = np.meshgrid(range(self.size[0]), range(self.size[1]), range(self.size[2]), sparse=False, indexing='ij')
+        self.xx = torch.tensor(xx, dtype=torch.float, device=self.device)
+        self.yy = torch.tensor(yy, dtype=torch.float, device=self.device)
+        self.zz = torch.tensor(zz, dtype=torch.float, device=self.device)
+        self.c = torch.tensor((np.array(self.size) - 1) / 2, dtype=torch.float, device=self.device)
+        self.xc = self.xx - self.c[0]
+        self.yc = self.yy - self.c[1]
+        self.zc = self.zz - self.c[2]
+        return
+    
+    def prepare_one_hot(self): 
+        if self.synth_args.left_hemis_only:
+            n_labels = len(label_list_segmentation_brainseg_left)
+            label_list_segmentation = label_list_segmentation_brainseg_left
+        else:
+            # Matrix for one-hot encoding (includes a lookup-table)
+            n_labels = len(label_list_segmentation_brainseg_with_extracerebral)
+            label_list_segmentation = label_list_segmentation_brainseg_with_extracerebral
+
+        self.lut = torch.zeros(10000, dtype=torch.long, device=self.device)
+        for l in range(n_labels):
+            self.lut[label_list_segmentation[l]] = l
+        self.onehotmatrix = torch.eye(n_labels, dtype=torch.float, device=self.device)
+        
+        # useless for left_hemis_only
+        nlat = int((n_labels - n_neutral_labels_brainseg_with_extracerebral) / 2.0)
+        self.vflip = np.concatenate([np.array(range(n_neutral_labels_brainseg_with_extracerebral)),
+                                np.array(range(n_neutral_labels_brainseg_with_extracerebral + nlat, n_labels)),
+                                np.array(range(n_neutral_labels_brainseg_with_extracerebral, n_neutral_labels_brainseg_with_extracerebral + nlat))])
+        return
+
+    
+    def random_affine_transform(self, shp):
+        rotations = (2 * self.synth_args.max_rotation * np.random.rand(3) - self.synth_args.max_rotation) / 180.0 * np.pi
+        shears = (2 * self.synth_args.max_shear * np.random.rand(3) - self.synth_args.max_shear)
+        scalings = 1 + (2 * self.synth_args.max_scaling * np.random.rand(3) - self.synth_args.max_scaling)
+        scaling_factor_distances = np.prod(scalings) ** .33333333333 
+        A = torch.tensor(make_affine_matrix(rotations, shears, scalings), dtype=torch.float, device=self.device)
+
+        # sample center
+        if self.synth_args.random_shift:
+            max_shift = (torch.tensor(np.array(shp[0:3]) - self.size, dtype=torch.float, device=self.device)) / 2
+            max_shift[max_shift < 0] = 0
+            c2 = torch.tensor((np.array(shp[0:3]) - 1)/2, dtype=torch.float, device=self.device) + (2 * (max_shift * torch.rand(3, dtype=float, device=self.device)) - max_shift)
+        else:
+            c2 = torch.tensor((np.array(shp[0:3]) - 1)/2, dtype=torch.float, device=self.device)
+        return scaling_factor_distances, A, c2
+
+    def random_nonlinear_transform(self, photo_mode, spac):
+        nonlin_scale = self.synth_args.nonlin_scale_min + np.random.rand(1) * (self.synth_args.nonlin_scale_max - self.synth_args.nonlin_scale_min)
+        size_F_small = np.round(nonlin_scale * np.array(self.size)).astype(int).tolist()
+        if photo_mode:
+            size_F_small[1] = np.round(self.size[1]/spac).astype(int)
+        nonlin_std = self.synth_args.nonlin_std_max * np.random.rand()
+        Fsmall = nonlin_std * torch.randn([*size_F_small, 3], dtype=torch.float, device=self.device)
+        F = myzoom_torch(Fsmall, np.array(self.size) / size_F_small)
+        if photo_mode:
+            F[:, :, :, 1] = 0
+
+        if 'surface' in self.tasks: # TODO need to integrate the non-linear deformation fields for inverse
+            steplength = 1.0 / (2.0 ** self.synth_args.n_steps_svf_integration)
+            Fsvf = F * steplength
+            for _ in range(self.synth_args.n_steps_svf_integration):
+                Fsvf += fast_3D_interp_torch(Fsvf, self.xx + Fsvf[:, :, :, 0], self.yy + Fsvf[:, :, :, 1], self.zz + Fsvf[:, :, :, 2], 'linear')
+            Fsvf_neg = -F * steplength
+            for _ in range(self.synth_args.n_steps_svf_integration):
+                Fsvf_neg += fast_3D_interp_torch(Fsvf_neg, self.xx + Fsvf_neg[:, :, :, 0], self.yy + Fsvf_neg[:, :, :, 1], self.zz + Fsvf_neg[:, :, :, 2], 'linear')
+            F = Fsvf
+            Fneg = Fsvf_neg
+        else:
+            Fneg = None
+        return F, Fneg
+    
+    def generate_deformation(self, setups, shp):
+
+        # generate affine deformation
+        scaling_factor_distances, A, c2 = self.random_affine_transform(shp)
+        
+        # generate nonlinear deformation 
+        if self.synth_args.nonlinear_transform:
+            F, Fneg = self.random_nonlinear_transform(setups['photo_mode'], setups['spac']) 
+        else:
+            F, Fneg = None, None
+
+        # deform the image grid 
+        xx2, yy2, zz2, x1, y1, z1, x2, y2, z2 = self.deform_grid(shp, A, c2, F)  
+
+        return {'scaling_factor_distances': scaling_factor_distances, 
+                'A': A, 
+                'c2': c2, 
+                'F': F, 
+                'Fneg': Fneg, 
+                'grid': [xx2, yy2, zz2, x1, y1, z1, x2, y2, z2], 
+                }
+
+
+    def get_left_hemis_mask(self, grid): 
+        [xx2, yy2, zz2, x1, y1, z1, x2, y2, z2] = grid
+
+        if self.synth_args.left_hemis_only: 
+            S, aff, res = read_image(self.modalities['segmentation']) # read seg map
+            S = torch.squeeze(torch.from_numpy(S.get_fdata()[x1:x2, y1:y2, z1:z2].astype(int))).to(self.device)
+            S = self.lut[S.int()] # mask out non-left labels
+            X, aff, res = read_image(self.modalities['registration'][0]) # read_mni_coord_X
+            X = torch.squeeze(torch.from_numpy(X.get_fdata()[x1:x2, y1:y2, z1:z2])).to(self.device)
+            self.hemis_mask = ((S > 0) & (X < 0)).int()
+        else:
+            self.hemis_mask = None
+    
+    def deform_grid(self, shp, A, c2, F): 
+        if F is not None:
+            # deform the images (we do nonlinear "first" ie after so we can do heavy coronal deformations in photo mode)
+            xx1 = self.xc + F[:, :, :, 0]
+            yy1 = self.yc + F[:, :, :, 1]
+            zz1 = self.zc + F[:, :, :, 2]
+        else:
+            xx1 = self.xc
+            yy1 = self.yc
+            zz1 = self.zc
+ 
+        xx2 = A[0, 0] * xx1 + A[0, 1] * yy1 + A[0, 2] * zz1 + c2[0]
+        yy2 = A[1, 0] * xx1 + A[1, 1] * yy1 + A[1, 2] * zz1 + c2[1]
+        zz2 = A[2, 0] * xx1 + A[2, 1] * yy1 + A[2, 2] * zz1 + c2[2]  
+        xx2[xx2 < 0] = 0
+        yy2[yy2 < 0] = 0
+        zz2[zz2 < 0] = 0
+        xx2[xx2 > (shp[0] - 1)] = shp[0] - 1
+        yy2[yy2 > (shp[1] - 1)] = shp[1] - 1
+        zz2[zz2 > (shp[2] - 1)] = shp[2] - 1
+
+        # Get the margins for reading images
+        x1 = torch.floor(torch.min(xx2))
+        y1 = torch.floor(torch.min(yy2))
+        z1 = torch.floor(torch.min(zz2))
+        x2 = 1+torch.ceil(torch.max(xx2))
+        y2 = 1 + torch.ceil(torch.max(yy2))
+        z2 = 1 + torch.ceil(torch.max(zz2))
+        xx2 -= x1
+        yy2 -= y1
+        zz2 -= z1
+
+        x1 = x1.cpu().numpy().astype(int)
+        y1 = y1.cpu().numpy().astype(int)
+        z1 = z1.cpu().numpy().astype(int)
+        x2 = x2.cpu().numpy().astype(int)
+        y2 = y2.cpu().numpy().astype(int)
+        z2 = z2.cpu().numpy().astype(int)
+
+        return xx2, yy2, zz2, x1, y1, z1, x2, y2, z2
+
+
+    def augment_sample(self, name, I_def, setups, deform_dict, res, target, pathol_direction = None, input_mode = 'synth'):
+
+        sample = {}
+        [xx2, yy2, zz2, x1, y1, z1, x2, y2, z2] = deform_dict['grid']
+
+        if not isinstance(I_def, torch.Tensor):
+            I_def = torch.squeeze(torch.tensor(I_def.get_fdata()[x1:x2, y1:y2, z1:z2].astype(float), dtype=torch.float, device=self.device))
+            if self.hemis_mask is not None:
+                I_def[self.hemis_mask == 0] = 0
+            # Deform grid
+            I_def = fast_3D_interp_torch(I_def, xx2, yy2, zz2, 'linear')
+
+        if input_mode == 'CT':
+            I_def = torch.clamp(I_def, min = 0., max = 80.)
+
+        if 'pathology' in target and isinstance(target['pathology'], torch.Tensor) and target['pathology'].sum() > 0:
+            I_def = self.encode_pathology(I_def, target['pathology'], target['pathology_prob'], pathol_direction)
+            I_def[I_def < 0.] = 0.
+        else: 
+            target['pathology'] = 0.
+            target['pathology_prob'] = 0.  
+
+        # Augment sample
+        aux_dict = {}
+        augmentation_steps = self.augmentation_steps['synth'] if input_mode == 'synth' else self.augmentation_steps['real']
+        for func_name in augmentation_steps:
+            I_def, aux_dict = augmentation_funcs[func_name](I = I_def, aux_dict = aux_dict, cfg = self.gen_args.generator, 
+                                                         input_mode = input_mode, setups = setups, size = self.size, res = res, device = self.device)
+
+
+        # Back to original resolution 
+        if self.synth_args.bspline_zooming:
+            I_def = interpol.resize(I_def, shape=self.size, anchor='edge', interpolation=3, bound='dct2', prefilter=True) 
+        else:
+            I_def = myzoom_torch(I_def, 1 / aux_dict['factors']) 
+            
+        maxi = torch.max(I_def)
+        I_final = I_def / maxi
+
+        if 'super_resolution' in self.tasks: 
+            SRresidual = aux_dict['high_res'] / maxi - I_final
+            sample.update({'high_res_residual': torch.flip(SRresidual, [0])[None] if setups['flip'] else SRresidual[None]})
+
+
+        sample.update({'input': torch.flip(I_final, [0])[None] if setups['flip'] else I_final[None]})
+        if 'bias_field' in self.tasks and input_mode != 'CT':
+            sample.update({'bias_field_log': torch.flip(aux_dict['BFlog'], [0])[None] if setups['flip'] else aux_dict['BFlog'][None]})
+
+        return sample 
+    
+
+    def generate_sample(self, name, G, setups, deform_dict, res, target):  
+        
+        [xx2, yy2, zz2, x1, y1, z1, x2, y2, z2] = deform_dict['grid']
+
+        # Generate contrasts
+        mus, sigmas = self.get_contrast(setups['photo_mode'])
+
+        G = torch.squeeze(torch.tensor(G.get_fdata()[x1:x2, y1:y2, z1:z2].astype(float), dtype=torch.float, device=self.device))
+        #G[G > 255] = 0 # kill extracerebral regions
+        G[G == 77] = 2 # merge WM lesion to white matter region
+        if self.hemis_mask is not None:
+            G[self.hemis_mask == 0] = 0
+        Gr = torch.round(G).long()
+        
+        SYN = mus[Gr] + sigmas[Gr] * torch.randn(Gr.shape, dtype=torch.float, device=self.device)
+        SYN[SYN < 0] = 0
+        #SYN /= mus[2] # normalize by WM
+        #SYN = gaussian_blur_3d(SYN, 0.5*np.ones(3), self.device) # cosmetic
+
+        SYN = fast_3D_interp_torch(SYN, xx2, yy2, zz2) 
+
+        # Make random linear combinations
+        if np.random.rand() < self.gen_args.mix_synth_prob: 
+            v = torch.rand(4)
+            v[2] = 0 if 'T2' not in self.modalities else v[2]
+            v[3] = 0 if 'FLAIR' not in self.modalities else v[3]
+            v /= torch.sum(v) 
+            SYN = v[0] * SYN + v[1] * target['T1'][0]
+            if 'T2' in self.modalities:
+                SYN += v[2] * target['T2'][0]
+            if 'FLAIR' in self.modalities:
+                SYN += v[3] * target['FLAIR'][0] 
+            
+        if 'pathology' in target and isinstance(target['pathology'], torch.Tensor) and target['pathology'].sum() > 0:
+            SYN_cerebral = SYN.clone()
+            SYN_cerebral[Gr == 0] = 0
+            SYN_cerebral = fast_3D_interp_torch(SYN_cerebral, xx2, yy2, zz2)[None]
+
+            wm_mask = (Gr==2) | (Gr==41)
+            wm_mean = (SYN * wm_mask).sum() / wm_mask.sum()  
+            gm_mask = (Gr!=0) & (Gr!=2) & (Gr!=41)
+            gm_mean = (SYN * gm_mask).sum() / gm_mask.sum()
+
+            target['pathology'][SYN_cerebral == 0] = 0
+            target['pathology_prob'][SYN_cerebral == 0] = 0 
+            # determine to be T1-resembled or T2-resembled
+            #if pathol_direction: lesion should be brigher than WM.mean() 
+            # pathol_direction: +1: T2-like; -1: T1-like
+            pathol_direction = self.get_pathology_direction('synth', gm_mean > wm_mean)
+        else:
+            pathol_direction = None
+            target['pathology'] = 0.
+            target['pathology_prob'] = 0. 
+            
+        SYN[SYN < 0.] = 0.
+        return target['pathology'], target['pathology_prob'], self.augment_sample(name, SYN, setups, deform_dict, res, target, pathol_direction = pathol_direction)
+    
+    def get_pathology_direction(self, input_mode, pathol_direction = None):  
+        #if np.random.rand() < 0.1: # in some (rare) cases, randomly pick the direction
+        #    return random.choice([True, False])
+        
+        if pathol_direction is not None: # for synth image
+            return pathol_direction
+        
+        if input_mode in ['T1', 'CT']:
+            return False
+        
+        if input_mode in ['T2', 'FLAIR']:
+            return True
+        
+        return random.choice([True, False])
+
+
+    def get_contrast(self, photo_mode):
+        # Sample Gaussian image
+        mus = 25 + 200 * torch.rand(256, dtype=torch.float, device=self.device)
+        sigmas = 5 + 20 * torch.rand(256, dtype=torch.float, device=self.device)
+
+        if np.random.rand() < self.synth_args.ct_prob:
+            darker = 25 + 10 * torch.rand(1, dtype=torch.float, device=self.device)[0]
+            for l in ct_brightness_group['darker']:
+                mus[l] = darker
+            dark = 90 + 20 * torch.rand(1, dtype=torch.float, device=self.device)[0]
+            for l in ct_brightness_group['dark']:
+                mus[l] = dark
+            bright = 110 + 20 * torch.rand(1, dtype=torch.float, device=self.device)[0]
+            for l in ct_brightness_group['bright']:
+                mus[l] = bright
+            brighter = 150 + 50 * torch.rand(1, dtype=torch.float, device=self.device)[0]
+            for l in ct_brightness_group['brighter']:
+                mus[l] = brighter
+                
+        if photo_mode or np.random.rand(1)<0.5: # set the background to zero every once in a while (or always in photo mode)
+            mus[0] = 0
+
+        #  partial volume
+        # 1 = lesion, 2 = WM, 3 = GM, 4 = CSF
+        v = 0.02 * torch.arange(50).to(self.device)
+        mus[100:150] = mus[1] * (1 - v) + mus[2] * v
+        mus[150:200] = mus[2] * (1 - v) + mus[3] * v
+        mus[200:250] = mus[3] * (1 - v) + mus[4] * v
+        mus[250] = mus[4]
+        sigmas[100:150] = torch.sqrt(sigmas[1]**2 * (1 - v) + sigmas[2]**2 * v)
+        sigmas[150:200] = torch.sqrt(sigmas[2]**2 * (1 - v) + sigmas[3]**2 * v)
+        sigmas[200:250] = torch.sqrt(sigmas[3]**2 * (1 - v) + sigmas[4]**2 * v)
+        sigmas[250] = sigmas[4]
+
+        return mus, sigmas
+    
+    def get_setup_params(self): 
+
+        if self.synth_args.left_hemis_only:
+            hemis = 'left'
+        else:
+            hemis = 'both' 
+
+        if self.synth_args.low_res_only:
+            photo_mode = False
+        elif self.synth_args.left_hemis_only:
+            photo_mode = True
+        else:
+            photo_mode = np.random.rand() < self.synth_args.photo_prob
+            
+        pathol_mode = np.random.rand() < self.synth_args.pathology_prob
+        pathol_random_shape = np.random.rand() < self.synth_args.random_shape_prob
+        spac = 2.5 + 10 * np.random.rand() if photo_mode else None  
+        flip = np.random.randn() < self.synth_args.flip_prob if not self.synth_args.left_hemis_only else False
+        
+        if photo_mode: 
+            resolution = np.array([self.res_training_data[0], spac, self.res_training_data[2]])
+            thickness = np.array([self.res_training_data[0], 0.1, self.res_training_data[2]])
+        else:
+            resolution, thickness = resolution_sampler(self.synth_args.low_res_only)
+        return {'resolution': resolution, 'thickness': thickness, 
+                'photo_mode': photo_mode, 'pathol_mode': pathol_mode, 
+                'pathol_random_shape': pathol_random_shape,
+                'spac': spac, 'flip': flip, 'hemis': hemis}
+    
+    
+    def encode_pathology(self, I, P, Pprob, pathol_direction = None):
+
+
+        if pathol_direction is None: # True: T2/FLAIR-resembled, False: T1-resembled
+            pathol_direction = random.choice([True, False])
+
+        P, Pprob = torch.squeeze(P), torch.squeeze(Pprob)
+        I_mu = (I * P).sum() / P.sum()
+
+        p_mask = torch.round(P).long()
+        #pth_mus = I_mu/4 + I_mu/2 * torch.rand(10000, dtype=torch.float, device=self.device)
+        pth_mus = 3*I_mu/4 + I_mu/4 * torch.rand(10000, dtype=torch.float, device=self.device) # enforce the pathology pattern harder!
+        pth_mus = pth_mus if pathol_direction else -pth_mus 
+        pth_sigmas = I_mu/4 * torch.rand(10000, dtype=torch.float, device=self.device)
+        I += Pprob * (pth_mus[p_mask] + pth_sigmas[p_mask] * torch.randn(p_mask.shape, dtype=torch.float, device=self.device))
+        I[I < 0] = 0
+
+        #print('encode', P.shape, P.mean()) 
+        #print('pre', I_mu) 
+        #I_mu = (I * P).sum() / P.sum()
+        #print('post', I_mu)
+
+        return I
+    
+    def get_info(self, t1):
+
+        t1dm = t1[:-7] + 'T1w.defacingmask.nii'
+        t2 = t1[:-7] + 'T2w.nii'
+        t2dm = t1[:-7] + 'T2w.defacingmask.nii'
+        flair = t1[:-7] + 'FLAIR.nii'
+        flairdm = t1[:-7] + 'FLAIR.defacingmask.nii'
+        ct = t1[:-7] + 'CT.nii'
+        ctdm = t1[:-7] + 'CT.defacingmask.nii'
+        generation_labels = t1[:-7] + 'generation_labels.nii' 
+        segmentation_labels = t1[:-7] + self.gen_args.segment_prefix + '.nii'
+        #brain_dist_map = t1[:-7] + 'brain_dist_map.nii'
+        lp_dist_map = t1[:-7] + 'lp_dist_map.nii'
+        rp_dist_map = t1[:-7] + 'rp_dist_map.nii'
+        lw_dist_map = t1[:-7] + 'lw_dist_map.nii'
+        rw_dist_map = t1[:-7] + 'rw_dist_map.nii'
+        mni_reg_x = t1[:-7] + 'mni_reg.x.nii'
+        mni_reg_y = t1[:-7] + 'mni_reg.y.nii'
+        mni_reg_z = t1[:-7] + 'mni_reg.z.nii'
+
+
+        self.modalities = {'T1': t1, 'Gen': generation_labels, 'segmentation': segmentation_labels,   
+                           'distance': [lp_dist_map, lw_dist_map, rp_dist_map, rw_dist_map],
+                           'registration': [mni_reg_x, mni_reg_y, mni_reg_z]}
+
+        if os.path.isfile(t1dm):
+            self.modalities.update({'T1_DM': t1dm}) 
+        if os.path.isfile(t2):
+            self.modalities.update({'T2': t2}) 
+        if os.path.isfile(t2dm):
+            self.modalities.update({'T2_DM': t2dm}) 
+        if os.path.isfile(flair):
+            self.modalities.update({'FLAIR': flair}) 
+        if os.path.isfile(flairdm):  
+            self.modalities.update({'FLAIR_DM': flairdm}) 
+        if os.path.isfile(ct): 
+            self.modalities.update({'CT': ct}) 
+        if os.path.isfile(ctdm): 
+            self.modalities.update({'CT_DM': ctdm}) 
+
+        return self.modalities
+
+
+    def read_input(self, idx):
+        """
+        determine input type according to prob (in generator/constants.py)
+        Logic: if np.random.rand() < real_image_prob and is real_image_exist --> input real images; otherwise, synthesize images. 
+        """
+        dataset_name, input_prob, t1_path, age = self.idx_to_path(idx)
+        case_name = os.path.basename(t1_path).split('.T1w.nii')[0]
+        self.modalities = self.get_info(t1_path)
+
+        prob = np.random.rand() 
+        if prob < input_prob['T1'] and 'T1' in self.modalities:
+            input_mode = 'T1'
+            img, aff, res = read_image(self.modalities['T1']) 
+        elif prob < input_prob['T2'] and 'T2' in self.modalities:
+            input_mode = 'T2'
+            img, aff, res = read_image(self.modalities['T2']) 
+        elif prob < input_prob['FLAIR'] and 'FLAIR' in self.modalities:
+            input_mode = 'FLAIR'
+            img, aff, res = read_image(self.modalities['FLAIR']) 
+        elif prob < input_prob['CT'] and 'CT' in self.modalities:
+            input_mode = 'CT'
+            img, aff, res = read_image(self.modalities['CT']) 
+        else:
+            input_mode = 'synth' 
+            img, aff, res = read_image(self.modalities['Gen']) 
+
+        return dataset_name, case_name, input_mode, img, aff, res, age
+    
+
+    def read_and_deform_target(self, idx, exist_keys, task_name, input_mode, setups, deform_dict, linear_weights = None):
+        current_target = {}
+        p_prob_path, augment, thres = None, False, 0.1
+
+        if task_name == 'pathology':
+            # NOTE: for now - encode pathology only for healthy cases
+            # TODO: what to do if the case has pathology itself? -- inconsistency between encoded pathol and the output
+            if self.pathology_type is None: # healthy
+                if setups['pathol_mode']: # and input_mode == 'synth':
+                    if setups['pathol_random_shape']:
+                        p_prob_path = 'random_shape'
+                        augment, thres = False, self.shape_gen_args.pathol_thres 
+                    else:
+                        p_prob_path = random.choice(pathology_prob_paths)
+                        augment, thres = self.synth_args.augment_pathology, self.shape_gen_args.pathol_thres 
+            else: 
+                pass
+                #p_prob_path = self.modalities['pathology_prob'] 
+
+            current_target = processing_funcs[task_name](exist_keys, task_name, p_prob_path, setups, deform_dict, self.device,
+                                                         mask = self.hemis_mask,
+                                                         augment = augment, 
+                                                         pde_func = self.adv_pde, 
+                                                         t = self.t, 
+                                                         shape_gen_args = self.shape_gen_args, 
+                                                         thres = thres
+                                                         )
+            
+        else:
+            if task_name in self.modalities:
+                current_target = processing_funcs[task_name](exist_keys, task_name, self.modalities[task_name], 
+                                                            setups, deform_dict, self.device, 
+                                                            mask = self.hemis_mask,
+                                                            cfg = self.gen_args, 
+                                                            onehotmatrix = self.onehotmatrix, 
+                                                            lut = self.lut, vflip = self.vflip
+                                                            )
+            else:
+                current_target = {task_name: 0.}
+        return current_target
+    
+        
+    def update_gen_args(self, new_args):
+        for key, value in vars(new_args).items():
+            vars(self.gen_args.generator)[key] = value 
+
+    def __getitem__(self, idx):
+        if torch.is_tensor(idx):
+            idx = idx.tolist()  
+
+        # read input: real or synthesized image, according to customized prob
+        dataset_name, case_name, input_mode, img, aff, res, age = self.read_input(idx)
+
+        # generate random values
+        setups = self.get_setup_params()
+
+        # sample random deformation
+        deform_dict = self.generate_deformation(setups, img.shape)
+
+        # get left_hemis_mask if needed
+        self.get_left_hemis_mask(deform_dict['grid'])
+
+        # read and deform target according to the assigned tasks
+        target = defaultdict(lambda: None)
+        target['name'] = case_name
+        target.update(self.read_and_deform_target(idx, target.keys(), 'T1', input_mode, setups, deform_dict))
+        target.update(self.read_and_deform_target(idx, target.keys(), 'T2', input_mode, setups, deform_dict)) 
+        target.update(self.read_and_deform_target(idx, target.keys(), 'FLAIR', input_mode, setups, deform_dict))
+        for task_name in self.tasks:
+            if task_name in processing_funcs.keys() and task_name not in ['T1', 'T2', 'FLAIR']: 
+                target.update(self.read_and_deform_target(idx, target.keys(), task_name, input_mode, setups, deform_dict))
+        
+
+        # process or generate input sample
+        if input_mode == 'synth':
+            self.update_gen_args(self.synth_image_args) # severe noise injection for real images
+            target['pathology'], target['pathology_prob'], sample = \
+                self.generate_sample(case_name, img, setups, deform_dict, res, target)  
+        else:
+            self.update_gen_args(self.real_image_args) # milder noise injection for real images
+            sample = self.augment_sample(case_name, img, setups, deform_dict, res, target,  
+                                        pathol_direction = self.get_pathology_direction(input_mode),input_mode = input_mode)
+
+        if setups['flip'] and isinstance(target['pathology'], torch.Tensor): # flipping should happen after P has been encoded
+            target['pathology'], target['pathology_prob'] = torch.flip(target['pathology'], [1]), torch.flip(target['pathology_prob'], [1]) 
+        
+        if age is not None:
+            target['age'] = age
+
+        return self.datasets_num, dataset_name, input_mode, target, sample
+
+
+
+
+# An example of customized dataset from BaseSynth
+class BrainIDGen(BaseGen):
+    """
+    BrainIDGen dataset
+    BrainIDGen enables intra-subject augmentation, i.e., each subject will have multiple augmentations
+    """
+    def __init__(self, gen_args, device='cpu'):  
+        super(BrainIDGen, self).__init__(gen_args, device)
+
+        self.all_samples = gen_args.generator.all_samples 
+        self.mild_samples = gen_args.generator.mild_samples 
+        self.mild_generator_args = gen_args.mild_generator
+        self.severe_generator_args = gen_args.severe_generator
+    
+    def __getitem__(self, idx):
+        if torch.is_tensor(idx):
+            idx = idx.tolist()  
+
+        # read input: real or synthesized image, according to customized prob 
+        dataset_name, case_name, input_mode, img, aff, res, age = self.read_input(idx)
+
+        # generate random values
+        setups = self.get_setup_params()
+
+        # sample random deformation
+        deform_dict = self.generate_deformation(setups, img.shape) 
+
+        # get left_hemis_mask if needed
+        self.get_left_hemis_mask(deform_dict['grid'])
+
+        # read and deform target according to the assigned tasks
+        target = defaultdict(lambda: 1.)
+        target['name'] = case_name
+        target.update(self.read_and_deform_target(idx, target.keys(), 'T1', input_mode, setups, deform_dict))
+        target.update(self.read_and_deform_target(idx, target.keys(), 'T2', input_mode, setups, deform_dict)) 
+        target.update(self.read_and_deform_target(idx, target.keys(), 'FLAIR', input_mode, setups, deform_dict))
+        for task_name in self.tasks:
+            if task_name in processing_funcs.keys() and task_name not in ['T1', 'T2', 'FLAIR']: 
+                target.update(self.read_and_deform_target(idx, target.keys(), task_name, input_mode, setups, deform_dict)) 
+
+        # process or generate intra-subject input samples 
+        samples = []
+        for i_sample in range(self.all_samples):
+            if i_sample < self.mild_samples:  
+                self.update_gen_args(self.mild_generator_args)
+                if input_mode == 'synth':
+                    self.update_gen_args(self.synth_image_args)
+                    target['pathology'], target['pathology_prob'], sample = \
+                        self.generate_sample(case_name, img, setups, deform_dict, res, target) 
+                else:
+                    self.update_gen_args(self.real_image_args)
+                    sample = self.augment_sample(case_name, img, setups, deform_dict, res, target,  
+                                                 pathol_direction = self.get_pathology_direction(input_mode),input_mode = input_mode)
+            else: 
+                self.update_gen_args(self.severe_generator_args)
+                if input_mode == 'synth':
+                    self.update_gen_args(self.synth_image_args)
+                    target['pathology'], target['pathology_prob'], sample = \
+                        self.generate_sample(case_name, img, setups, deform_dict, res, target)  
+                else:
+                    self.update_gen_args(self.real_image_args) 
+                    sample = self.augment_sample(case_name, img, setups, deform_dict, res, target, 
+                                                 pathol_direction = self.get_pathology_direction(input_mode),input_mode = input_mode)
+
+            samples.append(sample) 
+        
+        if setups['flip'] and isinstance(target['pathology'], torch.Tensor): # flipping should happen after P has been encoded
+            target['pathology'], target['pathology_prob'] = torch.flip(target['pathology'], [1]), torch.flip(target['pathology_prob'], [1]) 
+ 
+        if age is not None:
+            target['age'] = age
+        return self.datasets_num, dataset_name, input_mode, target, samples

Generator/interpol/__init__.py

@@ -0,0 +1,7 @@
+from .api import *
+from .resize import *
+from .restrict import *
+from . import backend
+
+from . import _version
+__version__ = _version.get_versions()['version']

Generator/interpol/_version.py

@@ -0,0 +1,623 @@
+
+# This file helps to compute a version number in source trees obtained from
+# git-archive tarball (such as those provided by githubs download-from-tag
+# feature). Distribution tarballs (built by setup.py sdist) and build
+# directories (produced by setup.py build) will contain a much shorter file
+# that just contains the computed version number.
+
+# This file is released into the public domain. Generated by
+# versioneer-0.20 (https://github.com/python-versioneer/python-versioneer)
+
+"""Git implementation of _version.py."""
+
+import errno
+import os
+import re
+import subprocess
+import sys
+
+
+def get_keywords():
+    """Get the keywords needed to look up the version information."""
+    # these strings will be replaced by git during git-archive.
+    # setup.py/versioneer.py will grep for the variable names, so they must
+    # each be defined on a line of their own. _version.py will just call
+    # get_keywords().
+    git_refnames = " (HEAD -> main, tag: 0.2.3)"
+    git_full = "414ed52c973b9d32e3e6a5a75c91cd5aab064f23"
+    git_date = "2023-04-17 20:36:50 -0400"
+    keywords = {"refnames": git_refnames, "full": git_full, "date": git_date}
+    return keywords
+
+
+class VersioneerConfig:  # pylint: disable=too-few-public-methods
+    """Container for Versioneer configuration parameters."""
+
+
+def get_config():
+    """Create, populate and return the VersioneerConfig() object."""
+    # these strings are filled in when 'setup.py versioneer' creates
+    # _version.py
+    cfg = VersioneerConfig()
+    cfg.VCS = "git"
+    cfg.style = "pep440"
+    cfg.tag_prefix = ""
+    cfg.parentdir_prefix = ""
+    cfg.versionfile_source = "interpol/_version.py"
+    cfg.verbose = False
+    return cfg
+
+
+class NotThisMethod(Exception):
+    """Exception raised if a method is not valid for the current scenario."""
+
+
+LONG_VERSION_PY = {}
+HANDLERS = {}
+
+
+def register_vcs_handler(vcs, method):  # decorator
+    """Create decorator to mark a method as the handler of a VCS."""
+    def decorate(f):
+        """Store f in HANDLERS[vcs][method]."""
+        if vcs not in HANDLERS:
+            HANDLERS[vcs] = {}
+        HANDLERS[vcs][method] = f
+        return f
+    return decorate
+
+
+# pylint:disable=too-many-arguments,consider-using-with # noqa
+def run_command(commands, args, cwd=None, verbose=False, hide_stderr=False,
+                env=None):
+    """Call the given command(s)."""
+    assert isinstance(commands, list)
+    process = None
+    for command in commands:
+        try:
+            dispcmd = str([command] + args)
+            # remember shell=False, so use git.cmd on windows, not just git
+            process = subprocess.Popen([command] + args, cwd=cwd, env=env,
+                                       stdout=subprocess.PIPE,
+                                       stderr=(subprocess.PIPE if hide_stderr
+                                               else None))
+            break
+        except EnvironmentError:
+            e = sys.exc_info()[1]
+            if e.errno == errno.ENOENT:
+                continue
+            if verbose:
+                print("unable to run %s" % dispcmd)
+                print(e)
+            return None, None
+    else:
+        if verbose:
+            print("unable to find command, tried %s" % (commands,))
+        return None, None
+    stdout = process.communicate()[0].strip().decode()
+    if process.returncode != 0:
+        if verbose:
+            print("unable to run %s (error)" % dispcmd)
+            print("stdout was %s" % stdout)
+        return None, process.returncode
+    return stdout, process.returncode
+
+
+def versions_from_parentdir(parentdir_prefix, root, verbose):
+    """Try to determine the version from the parent directory name.
+
+    Source tarballs conventionally unpack into a directory that includes both
+    the project name and a version string. We will also support searching up
+    two directory levels for an appropriately named parent directory
+    """
+    rootdirs = []
+
+    for _ in range(3):
+        dirname = os.path.basename(root)
+        if dirname.startswith(parentdir_prefix):
+            return {"version": dirname[len(parentdir_prefix):],
+                    "full-revisionid": None,
+                    "dirty": False, "error": None, "date": None}
+        rootdirs.append(root)
+        root = os.path.dirname(root)  # up a level
+
+    if verbose:
+        print("Tried directories %s but none started with prefix %s" %
+              (str(rootdirs), parentdir_prefix))
+    raise NotThisMethod("rootdir doesn't start with parentdir_prefix")
+
+
+@register_vcs_handler("git", "get_keywords")
+def git_get_keywords(versionfile_abs):
+    """Extract version information from the given file."""
+    # the code embedded in _version.py can just fetch the value of these
+    # keywords. When used from setup.py, we don't want to import _version.py,
+    # so we do it with a regexp instead. This function is not used from
+    # _version.py.
+    keywords = {}
+    try:
+        with open(versionfile_abs, "r") as fobj:
+            for line in fobj:
+                if line.strip().startswith("git_refnames ="):
+                    mo = re.search(r'=\s*"(.*)"', line)
+                    if mo:
+                        keywords["refnames"] = mo.group(1)
+                if line.strip().startswith("git_full ="):
+                    mo = re.search(r'=\s*"(.*)"', line)
+                    if mo:
+                        keywords["full"] = mo.group(1)
+                if line.strip().startswith("git_date ="):
+                    mo = re.search(r'=\s*"(.*)"', line)
+                    if mo:
+                        keywords["date"] = mo.group(1)
+    except EnvironmentError:
+        pass
+    return keywords
+
+
+@register_vcs_handler("git", "keywords")
+def git_versions_from_keywords(keywords, tag_prefix, verbose):
+    """Get version information from git keywords."""
+    if "refnames" not in keywords:
+        raise NotThisMethod("Short version file found")
+    date = keywords.get("date")
+    if date is not None:
+        # Use only the last line.  Previous lines may contain GPG signature
+        # information.
+        date = date.splitlines()[-1]
+
+        # git-2.2.0 added "%cI", which expands to an ISO-8601 -compliant
+        # datestamp. However we prefer "%ci" (which expands to an "ISO-8601
+        # -like" string, which we must then edit to make compliant), because
+        # it's been around since git-1.5.3, and it's too difficult to
+        # discover which version we're using, or to work around using an
+        # older one.
+        date = date.strip().replace(" ", "T", 1).replace(" ", "", 1)
+    refnames = keywords["refnames"].strip()
+    if refnames.startswith("$Format"):
+        if verbose:
+            print("keywords are unexpanded, not using")
+        raise NotThisMethod("unexpanded keywords, not a git-archive tarball")
+    refs = {r.strip() for r in refnames.strip("()").split(",")}
+    # starting in git-1.8.3, tags are listed as "tag: foo-1.0" instead of
+    # just "foo-1.0". If we see a "tag: " prefix, prefer those.
+    TAG = "tag: "
+    tags = {r[len(TAG):] for r in refs if r.startswith(TAG)}
+    if not tags:
+        # Either we're using git < 1.8.3, or there really are no tags. We use
+        # a heuristic: assume all version tags have a digit. The old git %d
+        # expansion behaves like git log --decorate=short and strips out the
+        # refs/heads/ and refs/tags/ prefixes that would let us distinguish
+        # between branches and tags. By ignoring refnames without digits, we
+        # filter out many common branch names like "release" and
+        # "stabilization", as well as "HEAD" and "master".
+        tags = {r for r in refs if re.search(r'\d', r)}
+        if verbose:
+            print("discarding '%s', no digits" % ",".join(refs - tags))
+    if verbose:
+        print("likely tags: %s" % ",".join(sorted(tags)))
+    for ref in sorted(tags):
+        # sorting will prefer e.g. "2.0" over "2.0rc1"
+        if ref.startswith(tag_prefix):
+            r = ref[len(tag_prefix):]
+            # Filter out refs that exactly match prefix or that don't start
+            # with a number once the prefix is stripped (mostly a concern
+            # when prefix is '')
+            if not re.match(r'\d', r):
+                continue
+            if verbose:
+                print("picking %s" % r)
+            return {"version": r,
+                    "full-revisionid": keywords["full"].strip(),
+                    "dirty": False, "error": None,
+                    "date": date}
+    # no suitable tags, so version is "0+unknown", but full hex is still there
+    if verbose:
+        print("no suitable tags, using unknown + full revision id")
+    return {"version": "0+unknown",
+            "full-revisionid": keywords["full"].strip(),
+            "dirty": False, "error": "no suitable tags", "date": None}
+
+
+@register_vcs_handler("git", "pieces_from_vcs")
+def git_pieces_from_vcs(tag_prefix, root, verbose, runner=run_command):
+    """Get version from 'git describe' in the root of the source tree.
+
+    This only gets called if the git-archive 'subst' keywords were *not*
+    expanded, and _version.py hasn't already been rewritten with a short
+    version string, meaning we're inside a checked out source tree.
+    """
+    GITS = ["git"]
+    if sys.platform == "win32":
+        GITS = ["git.cmd", "git.exe"]
+
+    _, rc = runner(GITS, ["rev-parse", "--git-dir"], cwd=root,
+                   hide_stderr=True)
+    if rc != 0:
+        if verbose:
+            print("Directory %s not under git control" % root)
+        raise NotThisMethod("'git rev-parse --git-dir' returned error")
+
+    # if there is a tag matching tag_prefix, this yields TAG-NUM-gHEX[-dirty]
+    # if there isn't one, this yields HEX[-dirty] (no NUM)
+    describe_out, rc = runner(GITS, ["describe", "--tags", "--dirty",
+                                     "--always", "--long",
+                                     "--match", "%s*" % tag_prefix],
+                              cwd=root)
+    # --long was added in git-1.5.5
+    if describe_out is None:
+        raise NotThisMethod("'git describe' failed")
+    describe_out = describe_out.strip()
+    full_out, rc = runner(GITS, ["rev-parse", "HEAD"], cwd=root)
+    if full_out is None:
+        raise NotThisMethod("'git rev-parse' failed")
+    full_out = full_out.strip()
+
+    pieces = {}
+    pieces["long"] = full_out
+    pieces["short"] = full_out[:7]  # maybe improved later
+    pieces["error"] = None
+
+    branch_name, rc = runner(GITS, ["rev-parse", "--abbrev-ref", "HEAD"],
+                             cwd=root)
+    # --abbrev-ref was added in git-1.6.3
+    if rc != 0 or branch_name is None:
+        raise NotThisMethod("'git rev-parse --abbrev-ref' returned error")
+    branch_name = branch_name.strip()
+
+    if branch_name == "HEAD":
+        # If we aren't exactly on a branch, pick a branch which represents
+        # the current commit. If all else fails, we are on a branchless
+        # commit.
+        branches, rc = runner(GITS, ["branch", "--contains"], cwd=root)
+        # --contains was added in git-1.5.4
+        if rc != 0 or branches is None:
+            raise NotThisMethod("'git branch --contains' returned error")
+        branches = branches.split("\n")
+
+        # Remove the first line if we're running detached
+        if "(" in branches[0]:
+            branches.pop(0)
+
+        # Strip off the leading "* " from the list of branches.
+        branches = [branch[2:] for branch in branches]
+        if "master" in branches:
+            branch_name = "master"
+        elif not branches:
+            branch_name = None
+        else:
+            # Pick the first branch that is returned. Good or bad.
+            branch_name = branches[0]
+
+    pieces["branch"] = branch_name
+
+    # parse describe_out. It will be like TAG-NUM-gHEX[-dirty] or HEX[-dirty]
+    # TAG might have hyphens.
+    git_describe = describe_out
+
+    # look for -dirty suffix
+    dirty = git_describe.endswith("-dirty")
+    pieces["dirty"] = dirty
+    if dirty:
+        git_describe = git_describe[:git_describe.rindex("-dirty")]
+
+    # now we have TAG-NUM-gHEX or HEX
+
+    if "-" in git_describe:
+        # TAG-NUM-gHEX
+        mo = re.search(r'^(.+)-(\d+)-g([0-9a-f]+)$', git_describe)
+        if not mo:
+            # unparseable. Maybe git-describe is misbehaving?
+            pieces["error"] = ("unable to parse git-describe output: '%s'"
+                               % describe_out)
+            return pieces
+
+        # tag
+        full_tag = mo.group(1)
+        if not full_tag.startswith(tag_prefix):
+            if verbose:
+                fmt = "tag '%s' doesn't start with prefix '%s'"
+                print(fmt % (full_tag, tag_prefix))
+            pieces["error"] = ("tag '%s' doesn't start with prefix '%s'"
+                               % (full_tag, tag_prefix))
+            return pieces
+        pieces["closest-tag"] = full_tag[len(tag_prefix):]
+
+        # distance: number of commits since tag
+        pieces["distance"] = int(mo.group(2))
+
+        # commit: short hex revision ID
+        pieces["short"] = mo.group(3)
+
+    else:
+        # HEX: no tags
+        pieces["closest-tag"] = None
+        count_out, rc = runner(GITS, ["rev-list", "HEAD", "--count"], cwd=root)
+        pieces["distance"] = int(count_out)  # total number of commits
+
+    # commit date: see ISO-8601 comment in git_versions_from_keywords()
+    date = runner(GITS, ["show", "-s", "--format=%ci", "HEAD"], cwd=root)[0].strip()
+    # Use only the last line.  Previous lines may contain GPG signature
+    # information.
+    date = date.splitlines()[-1]
+    pieces["date"] = date.strip().replace(" ", "T", 1).replace(" ", "", 1)
+
+    return pieces
+
+
+def plus_or_dot(pieces):
+    """Return a + if we don't already have one, else return a ."""
+    if "+" in pieces.get("closest-tag", ""):
+        return "."
+    return "+"
+
+
+def render_pep440(pieces):
+    """Build up version string, with post-release "local version identifier".
+
+    Our goal: TAG[+DISTANCE.gHEX[.dirty]] . Note that if you
+    get a tagged build and then dirty it, you'll get TAG+0.gHEX.dirty
+
+    Exceptions:
+    1: no tags. git_describe was just HEX. 0+untagged.DISTANCE.gHEX[.dirty]
+    """
+    if pieces["closest-tag"]:
+        rendered = pieces["closest-tag"]
+        if pieces["distance"] or pieces["dirty"]:
+            rendered += plus_or_dot(pieces)
+            rendered += "%d.g%s" % (pieces["distance"], pieces["short"])
+            if pieces["dirty"]:
+                rendered += ".dirty"
+    else:
+        # exception #1
+        rendered = "0+untagged.%d.g%s" % (pieces["distance"],
+                                          pieces["short"])
+        if pieces["dirty"]:
+            rendered += ".dirty"
+    return rendered
+
+
+def render_pep440_branch(pieces):
+    """TAG[[.dev0]+DISTANCE.gHEX[.dirty]] .
+
+    The ".dev0" means not master branch. Note that .dev0 sorts backwards
+    (a feature branch will appear "older" than the master branch).
+
+    Exceptions:
+    1: no tags. 0[.dev0]+untagged.DISTANCE.gHEX[.dirty]
+    """
+    if pieces["closest-tag"]:
+        rendered = pieces["closest-tag"]
+        if pieces["distance"] or pieces["dirty"]:
+            if pieces["branch"] != "master":
+                rendered += ".dev0"
+            rendered += plus_or_dot(pieces)
+            rendered += "%d.g%s" % (pieces["distance"], pieces["short"])
+            if pieces["dirty"]:
+                rendered += ".dirty"
+    else:
+        # exception #1
+        rendered = "0"
+        if pieces["branch"] != "master":
+            rendered += ".dev0"
+        rendered += "+untagged.%d.g%s" % (pieces["distance"],
+                                          pieces["short"])
+        if pieces["dirty"]:
+            rendered += ".dirty"
+    return rendered
+
+
+def render_pep440_pre(pieces):
+    """TAG[.post0.devDISTANCE] -- No -dirty.
+
+    Exceptions:
+    1: no tags. 0.post0.devDISTANCE
+    """
+    if pieces["closest-tag"]:
+        rendered = pieces["closest-tag"]
+        if pieces["distance"]:
+            rendered += ".post0.dev%d" % pieces["distance"]
+    else:
+        # exception #1
+        rendered = "0.post0.dev%d" % pieces["distance"]
+    return rendered
+
+
+def render_pep440_post(pieces):
+    """TAG[.postDISTANCE[.dev0]+gHEX] .
+
+    The ".dev0" means dirty. Note that .dev0 sorts backwards
+    (a dirty tree will appear "older" than the corresponding clean one),
+    but you shouldn't be releasing software with -dirty anyways.
+
+    Exceptions:
+    1: no tags. 0.postDISTANCE[.dev0]
+    """
+    if pieces["closest-tag"]:
+        rendered = pieces["closest-tag"]
+        if pieces["distance"] or pieces["dirty"]:
+            rendered += ".post%d" % pieces["distance"]
+            if pieces["dirty"]:
+                rendered += ".dev0"
+            rendered += plus_or_dot(pieces)
+            rendered += "g%s" % pieces["short"]
+    else:
+        # exception #1
+        rendered = "0.post%d" % pieces["distance"]
+        if pieces["dirty"]:
+            rendered += ".dev0"
+        rendered += "+g%s" % pieces["short"]
+    return rendered
+
+
+def render_pep440_post_branch(pieces):
+    """TAG[.postDISTANCE[.dev0]+gHEX[.dirty]] .
+
+    The ".dev0" means not master branch.
+
+    Exceptions:
+    1: no tags. 0.postDISTANCE[.dev0]+gHEX[.dirty]
+    """
+    if pieces["closest-tag"]:
+        rendered = pieces["closest-tag"]
+        if pieces["distance"] or pieces["dirty"]:
+            rendered += ".post%d" % pieces["distance"]
+            if pieces["branch"] != "master":
+                rendered += ".dev0"
+            rendered += plus_or_dot(pieces)
+            rendered += "g%s" % pieces["short"]
+            if pieces["dirty"]:
+                rendered += ".dirty"
+    else:
+        # exception #1
+        rendered = "0.post%d" % pieces["distance"]
+        if pieces["branch"] != "master":
+            rendered += ".dev0"
+        rendered += "+g%s" % pieces["short"]
+        if pieces["dirty"]:
+            rendered += ".dirty"
+    return rendered
+
+
+def render_pep440_old(pieces):
+    """TAG[.postDISTANCE[.dev0]] .
+
+    The ".dev0" means dirty.
+
+    Exceptions:
+    1: no tags. 0.postDISTANCE[.dev0]
+    """
+    if pieces["closest-tag"]:
+        rendered = pieces["closest-tag"]
+        if pieces["distance"] or pieces["dirty"]:
+            rendered += ".post%d" % pieces["distance"]
+            if pieces["dirty"]:
+                rendered += ".dev0"
+    else:
+        # exception #1
+        rendered = "0.post%d" % pieces["distance"]
+        if pieces["dirty"]:
+            rendered += ".dev0"
+    return rendered
+
+
+def render_git_describe(pieces):
+    """TAG[-DISTANCE-gHEX][-dirty].
+
+    Like 'git describe --tags --dirty --always'.
+
+    Exceptions:
+    1: no tags. HEX[-dirty]  (note: no 'g' prefix)
+    """
+    if pieces["closest-tag"]:
+        rendered = pieces["closest-tag"]
+        if pieces["distance"]:
+            rendered += "-%d-g%s" % (pieces["distance"], pieces["short"])
+    else:
+        # exception #1
+        rendered = pieces["short"]
+    if pieces["dirty"]:
+        rendered += "-dirty"
+    return rendered
+
+
+def render_git_describe_long(pieces):
+    """TAG-DISTANCE-gHEX[-dirty].
+
+    Like 'git describe --tags --dirty --always -long'.
+    The distance/hash is unconditional.
+
+    Exceptions:
+    1: no tags. HEX[-dirty]  (note: no 'g' prefix)
+    """
+    if pieces["closest-tag"]:
+        rendered = pieces["closest-tag"]
+        rendered += "-%d-g%s" % (pieces["distance"], pieces["short"])
+    else:
+        # exception #1
+        rendered = pieces["short"]
+    if pieces["dirty"]:
+        rendered += "-dirty"
+    return rendered
+
+
+def render(pieces, style):
+    """Render the given version pieces into the requested style."""
+    if pieces["error"]:
+        return {"version": "unknown",
+                "full-revisionid": pieces.get("long"),
+                "dirty": None,
+                "error": pieces["error"],
+                "date": None}
+
+    if not style or style == "default":
+        style = "pep440"  # the default
+
+    if style == "pep440":
+        rendered = render_pep440(pieces)
+    elif style == "pep440-branch":
+        rendered = render_pep440_branch(pieces)
+    elif style == "pep440-pre":
+        rendered = render_pep440_pre(pieces)
+    elif style == "pep440-post":
+        rendered = render_pep440_post(pieces)
+    elif style == "pep440-post-branch":
+        rendered = render_pep440_post_branch(pieces)
+    elif style == "pep440-old":
+        rendered = render_pep440_old(pieces)
+    elif style == "git-describe":
+        rendered = render_git_describe(pieces)
+    elif style == "git-describe-long":
+        rendered = render_git_describe_long(pieces)
+    else:
+        raise ValueError("unknown style '%s'" % style)
+
+    return {"version": rendered, "full-revisionid": pieces["long"],
+            "dirty": pieces["dirty"], "error": None,
+            "date": pieces.get("date")}
+
+
+def get_versions():
+    """Get version information or return default if unable to do so."""
+    # I am in _version.py, which lives at ROOT/VERSIONFILE_SOURCE. If we have
+    # __file__, we can work backwards from there to the root. Some
+    # py2exe/bbfreeze/non-CPython implementations don't do __file__, in which
+    # case we can only use expanded keywords.
+
+    cfg = get_config()
+    verbose = cfg.verbose
+
+    try:
+        return git_versions_from_keywords(get_keywords(), cfg.tag_prefix,
+                                          verbose)
+    except NotThisMethod:
+        pass
+
+    try:
+        root = os.path.realpath(__file__)
+        # versionfile_source is the relative path from the top of the source
+        # tree (where the .git directory might live) to this file. Invert
+        # this to find the root from __file__.
+        for _ in cfg.versionfile_source.split('/'):
+            root = os.path.dirname(root)
+    except NameError:
+        return {"version": "0+unknown", "full-revisionid": None,
+                "dirty": None,
+                "error": "unable to find root of source tree",
+                "date": None}
+
+    try:
+        pieces = git_pieces_from_vcs(cfg.tag_prefix, root, verbose)
+        return render(pieces, cfg.style)
+    except NotThisMethod:
+        pass
+
+    try:
+        if cfg.parentdir_prefix:
+            return versions_from_parentdir(cfg.parentdir_prefix, root, verbose)
+    except NotThisMethod:
+        pass
+
+    return {"version": "0+unknown", "full-revisionid": None,
+            "dirty": None,
+            "error": "unable to compute version", "date": None}

Generator/interpol/api.py

@@ -0,0 +1,560 @@
+"""High level interpolation API"""
+
+__all__ = ['grid_pull', 'grid_push', 'grid_count', 'grid_grad',
+           'spline_coeff', 'spline_coeff_nd',
+           'identity_grid', 'add_identity_grid', 'add_identity_grid_']
+
+import torch
+from .utils import expanded_shape, matvec
+from .jit_utils import movedim1, meshgrid
+from .autograd import (GridPull, GridPush, GridCount, GridGrad,
+                       SplineCoeff, SplineCoeffND)
+from . import backend, jitfields
+
+_doc_interpolation = \
+"""`interpolation` can be an int, a string or an InterpolationType.
+    Possible values are:
+        - 0 or 'nearest'
+        - 1 or 'linear'
+        - 2 or 'quadratic'
+        - 3 or 'cubic'
+        - 4 or 'fourth'
+        - 5 or 'fifth'
+        - etc.
+    A list of values can be provided, in the order [W, H, D],
+    to specify dimension-specific interpolation orders."""
+
+_doc_bound = \
+"""`bound` can be an int, a string or a BoundType.
+    Possible values are:
+        - 'replicate'  or 'nearest'     :  a  a  a  |  a  b  c  d  |  d  d  d
+        - 'dct1'       or 'mirror'      :  d  c  b  |  a  b  c  d  |  c  b  a
+        - 'dct2'       or 'reflect'     :  c  b  a  |  a  b  c  d  |  d  c  b
+        - 'dst1'       or 'antimirror'  : -b -a  0  |  a  b  c  d  |  0 -d -c
+        - 'dst2'       or 'antireflect' : -c -b -a  |  a  b  c  d  | -d -c -b
+        - 'dft'        or 'wrap'        :  b  c  d  |  a  b  c  d  |  a  b  c
+        - 'zero'       or 'zeros'       :  0  0  0  |  a  b  c  d  |  0  0  0
+    A list of values can be provided, in the order [W, H, D],
+    to specify dimension-specific boundary conditions.
+    Note that
+    - `dft` corresponds to circular padding
+    - `dct2` corresponds to Neumann boundary conditions (symmetric)
+    - `dst2` corresponds to Dirichlet boundary conditions (antisymmetric)
+    See https://en.wikipedia.org/wiki/Discrete_cosine_transform
+        https://en.wikipedia.org/wiki/Discrete_sine_transform"""
+
+_doc_bound_coeff = \
+"""`bound` can be an int, a string or a BoundType. 
+    Possible values are:
+        - 'replicate'  or 'nearest'     :  a  a  a  |  a  b  c  d  |  d  d  d
+        - 'dct1'       or 'mirror'      :  d  c  b  |  a  b  c  d  |  c  b  a
+        - 'dct2'       or 'reflect'     :  c  b  a  |  a  b  c  d  |  d  c  b
+        - 'dst1'       or 'antimirror'  : -b -a  0  |  a  b  c  d  |  0 -d -c
+        - 'dst2'       or 'antireflect' : -c -b -a  |  a  b  c  d  | -d -c -b
+        - 'dft'        or 'wrap'        :  b  c  d  |  a  b  c  d  |  a  b  c
+        - 'zero'       or 'zeros'       :  0  0  0  |  a  b  c  d  |  0  0  0
+    A list of values can be provided, in the order [W, H, D],
+    to specify dimension-specific boundary conditions.
+    Note that
+    - `dft` corresponds to circular padding
+    - `dct1` corresponds to mirroring about the center of the first/last voxel
+    - `dct2` corresponds to mirroring about the edge of the first/last voxel
+    See https://en.wikipedia.org/wiki/Discrete_cosine_transform
+        https://en.wikipedia.org/wiki/Discrete_sine_transform
+        
+    /!\ Only 'dct1', 'dct2' and 'dft' are implemented for interpolation
+        orders >= 6."""
+
+_ref_coeff = \
+"""..[1]  M. Unser, A. Aldroubi and M. Eden.
+       "B-Spline Signal Processing: Part I-Theory,"
+       IEEE Transactions on Signal Processing 41(2):821-832 (1993).
+..[2]  M. Unser, A. Aldroubi and M. Eden.
+       "B-Spline Signal Processing: Part II-Efficient Design and Applications,"
+       IEEE Transactions on Signal Processing 41(2):834-848 (1993).
+..[3]  M. Unser.
+       "Splines: A Perfect Fit for Signal and Image Processing,"
+       IEEE Signal Processing Magazine 16(6):22-38 (1999).
+"""
+
+
+def _preproc(grid, input=None, mode=None):
+    """Preprocess tensors for pull/push/count/grad
+
+    Low level bindings expect inputs of shape
+    [batch, channel, *spatial] and [batch, *spatial, dim], whereas
+    the high level python API accepts inputs of shape
+    [..., [channel], *spatial] and [..., *spatial, dim].
+
+    This function broadcasts and reshapes the input tensors accordingly.
+            /!\\ This *can* trigger large allocations /!\\
+    """
+    dim = grid.shape[-1]
+    if input is None:
+        spatial = grid.shape[-dim-1:-1]
+        batch = grid.shape[:-dim-1]
+        grid = grid.reshape([-1, *spatial, dim])
+        info = dict(batch=batch, channel=[1] if batch else [], dim=dim)
+        return grid, info
+
+    grid_spatial = grid.shape[-dim-1:-1]
+    grid_batch = grid.shape[:-dim-1]
+    input_spatial = input.shape[-dim:]
+    channel = 0 if input.dim() == dim else input.shape[-dim-1]
+    input_batch = input.shape[:-dim-1]
+
+    if mode == 'push':
+        grid_spatial = input_spatial = expanded_shape(grid_spatial, input_spatial)
+
+    # broadcast and reshape
+    batch = expanded_shape(grid_batch, input_batch)
+    grid = grid.expand([*batch, *grid_spatial, dim])
+    grid = grid.reshape([-1, *grid_spatial, dim])
+    input = input.expand([*batch, channel or 1, *input_spatial])
+    input = input.reshape([-1, channel or 1, *input_spatial])
+
+    out_channel = [channel] if channel else ([1] if batch else [])
+    info = dict(batch=batch, channel=out_channel, dim=dim)
+    return grid, input, info
+
+
+def _postproc(out, shape_info, mode):
+    """Postprocess tensors for pull/push/count/grad"""
+    dim = shape_info['dim']
+    if mode != 'grad':
+        spatial = out.shape[-dim:]
+        feat = []
+    else:
+        spatial = out.shape[-dim-1:-1]
+        feat = [out.shape[-1]]
+    batch = shape_info['batch']
+    channel = shape_info['channel']
+
+    out = out.reshape([*batch, *channel, *spatial, *feat])
+    return out
+
+
+def grid_pull(input, grid, interpolation='linear', bound='zero',
+              extrapolate=False, prefilter=False):
+    """Sample an image with respect to a deformation field.
+
+    Notes
+    -----
+    {interpolation}
+
+    {bound}
+    
+    If the input dtype is not a floating point type, the input image is 
+    assumed to contain labels. Then, unique labels are extracted 
+    and resampled individually, making them soft labels. Finally, 
+    the label map is reconstructed from the individual soft labels by 
+    assigning the label with maximum soft value.
+
+    Parameters
+    ----------
+    input : (..., [channel], *inshape) tensor
+        Input image.
+    grid : (..., *outshape, dim) tensor
+        Transformation field.
+    interpolation : int or sequence[int], default=1
+        Interpolation order.
+    bound : BoundType or sequence[BoundType], default='zero'
+        Boundary conditions.
+    extrapolate : bool or int, default=True
+        Extrapolate out-of-bound data.
+    prefilter : bool, default=False
+        Apply spline pre-filter (= interpolates the input)
+
+    Returns
+    -------
+    output : (..., [channel], *outshape) tensor
+        Deformed image.
+
+    """
+    if backend.jitfields and jitfields.available:
+        return jitfields.grid_pull(input, grid, interpolation, bound,
+                                   extrapolate, prefilter)
+
+    grid, input, shape_info = _preproc(grid, input)
+    batch, channel = input.shape[:2]
+    dim = grid.shape[-1]
+
+    if not input.dtype.is_floating_point:
+        # label map -> specific processing
+        out = input.new_zeros([batch, channel, *grid.shape[1:-1]])
+        pmax = grid.new_zeros([batch, channel, *grid.shape[1:-1]])
+        for label in input.unique():
+            soft = (input == label).to(grid.dtype)
+            if prefilter:
+                input = spline_coeff_nd(soft, interpolation=interpolation,
+                                        bound=bound, dim=dim, inplace=True)
+            soft = GridPull.apply(soft, grid, interpolation, bound, extrapolate)
+            out[soft > pmax] = label
+            pmax = torch.max(pmax, soft)
+    else:
+        if prefilter:
+            input = spline_coeff_nd(input, interpolation=interpolation,
+                                    bound=bound, dim=dim)
+        out = GridPull.apply(input, grid, interpolation, bound, extrapolate)
+
+    return _postproc(out, shape_info, mode='pull')
+
+
+def grid_push(input, grid, shape=None, interpolation='linear', bound='zero',
+              extrapolate=False, prefilter=False):
+    """Splat an image with respect to a deformation field (pull adjoint).
+
+    Notes
+    -----
+    {interpolation}
+
+    {bound}
+
+    Parameters
+    ----------
+    input : (..., [channel], *inshape) tensor
+        Input image.
+    grid : (..., *inshape, dim) tensor
+        Transformation field.
+    shape : sequence[int], default=inshape
+        Output shape
+    interpolation : int or sequence[int], default=1
+        Interpolation order.
+    bound : BoundType, or sequence[BoundType], default='zero'
+        Boundary conditions.
+    extrapolate : bool or int, default=True
+        Extrapolate out-of-bound data.
+    prefilter : bool, default=False
+        Apply spline pre-filter.
+
+    Returns
+    -------
+    output : (..., [channel], *shape) tensor
+        Spatted image.
+
+    """
+    if backend.jitfields and jitfields.available:
+        return jitfields.grid_push(input, grid, shape, interpolation, bound,
+                                   extrapolate, prefilter)
+
+    grid, input, shape_info = _preproc(grid, input, mode='push')
+    dim = grid.shape[-1]
+
+    if shape is None:
+        shape = tuple(input.shape[2:])
+
+    out = GridPush.apply(input, grid, shape, interpolation, bound, extrapolate)
+    if prefilter:
+        out = spline_coeff_nd(out, interpolation=interpolation, bound=bound,
+                              dim=dim, inplace=True)
+    return _postproc(out, shape_info, mode='push')
+
+
+def grid_count(grid, shape=None, interpolation='linear', bound='zero',
+               extrapolate=False):
+    """Splatting weights with respect to a deformation field (pull adjoint).
+
+    Notes
+    -----
+    {interpolation}
+
+    {bound}
+
+    Parameters
+    ----------
+    grid : (..., *inshape, dim) tensor
+        Transformation field.
+    shape : sequence[int], default=inshape
+        Output shape
+    interpolation : int or sequence[int], default=1
+        Interpolation order.
+    bound : BoundType, or sequence[BoundType], default='zero'
+        Boundary conditions.
+    extrapolate : bool or int, default=True
+        Extrapolate out-of-bound data.
+
+    Returns
+    -------
+    output : (..., [1], *shape) tensor
+        Splatted weights.
+
+    """
+    if backend.jitfields and jitfields.available:
+        return jitfields.grid_count(grid, shape, interpolation, bound, extrapolate)
+
+    grid, shape_info = _preproc(grid)
+    out = GridCount.apply(grid, shape, interpolation, bound, extrapolate)
+    return _postproc(out, shape_info, mode='count')
+
+
+def grid_grad(input, grid, interpolation='linear', bound='zero',
+              extrapolate=False, prefilter=False):
+    """Sample spatial gradients of an image with respect to a deformation field.
+    
+    Notes
+    -----
+    {interpolation}
+
+    {bound}
+
+    Parameters
+    ----------
+    input : (..., [channel], *inshape) tensor
+        Input image.
+    grid : (..., *inshape, dim) tensor
+        Transformation field.
+    shape : sequence[int], default=inshape
+        Output shape
+    interpolation : int or sequence[int], default=1
+        Interpolation order.
+    bound : BoundType, or sequence[BoundType], default='zero'
+        Boundary conditions.
+    extrapolate : bool or int, default=True
+        Extrapolate out-of-bound data.
+    prefilter : bool, default=False
+        Apply spline pre-filter (= interpolates the input)
+
+    Returns
+    -------
+    output : (..., [channel], *shape, dim) tensor
+        Sampled gradients.
+
+    """
+    if backend.jitfields and jitfields.available:
+        return jitfields.grid_grad(input, grid, interpolation, bound,
+                                   extrapolate, prefilter)
+
+    grid, input, shape_info = _preproc(grid, input)
+    dim = grid.shape[-1]
+    if prefilter:
+        input = spline_coeff_nd(input, interpolation, bound, dim)
+    out = GridGrad.apply(input, grid, interpolation, bound, extrapolate)
+    return _postproc(out, shape_info, mode='grad')
+
+
+def spline_coeff(input, interpolation='linear', bound='dct2', dim=-1,
+                 inplace=False):
+    """Compute the interpolating spline coefficients, for a given spline order
+    and boundary conditions, along a single dimension.
+
+    Notes
+    -----
+    {interpolation}
+
+    {bound}
+
+    References
+    ----------
+    {ref}
+
+
+    Parameters
+    ----------
+    input : tensor
+        Input image.
+    interpolation : int or sequence[int], default=1
+        Interpolation order.
+    bound : BoundType or sequence[BoundType], default='dct1'
+        Boundary conditions.
+    dim : int, default=-1
+        Dimension along which to process
+    inplace : bool, default=False
+        Process the volume in place.
+
+    Returns
+    -------
+    output : tensor
+        Coefficient image.
+
+    """
+    # This implementation is based on the file bsplines.c in SPM12, written
+    # by John Ashburner, which is itself based on the file coeff.c,
+    # written by Philippe Thevenaz: http://bigwww.epfl.ch/thevenaz/interpolation
+    # . DCT1 boundary conditions were derived by Thevenaz and Unser.
+    # . DFT boundary conditions were derived by John Ashburner.
+    # SPM12 is released under the GNU-GPL v2 license.
+    # Philippe Thevenaz's code does not have an explicit license as far
+    # as we know.
+    if backend.jitfields and jitfields.available:
+        return jitfields.spline_coeff(input, interpolation, bound,
+                                      dim, inplace)
+
+    out = SplineCoeff.apply(input, bound, interpolation, dim, inplace)
+    return out
+
+
+def spline_coeff_nd(input, interpolation='linear', bound='dct2', dim=None,
+                    inplace=False):
+    """Compute the interpolating spline coefficients, for a given spline order
+    and boundary conditions, along the last `dim` dimensions.
+
+    Notes
+    -----
+    {interpolation}
+
+    {bound}
+
+    References
+    ----------
+    {ref}
+
+    Parameters
+    ----------
+    input : (..., *spatial) tensor
+        Input image.
+    interpolation : int or sequence[int], default=1
+        Interpolation order.
+    bound : BoundType or sequence[BoundType], default='dct1'
+        Boundary conditions.
+    dim : int, default=-1
+        Number of spatial dimensions
+    inplace : bool, default=False
+        Process the volume in place.
+
+    Returns
+    -------
+    output : (..., *spatial) tensor
+        Coefficient image.
+
+    """
+    # This implementation is based on the file bsplines.c in SPM12, written
+    # by John Ashburner, which is itself based on the file coeff.c,
+    # written by Philippe Thevenaz: http://bigwww.epfl.ch/thevenaz/interpolation
+    # . DCT1 boundary conditions were derived by Thevenaz and Unser.
+    # . DFT boundary conditions were derived by John Ashburner.
+    # SPM12 is released under the GNU-GPL v2 license.
+    # Philippe Thevenaz's code does not have an explicit license as far
+    # as we know.
+    if backend.jitfields and jitfields.available:
+        return jitfields.spline_coeff_nd(input, interpolation, bound,
+                                         dim, inplace)
+
+    out = SplineCoeffND.apply(input, bound, interpolation, dim, inplace)
+    return out
+
+
+grid_pull.__doc__ = grid_pull.__doc__.format(
+    interpolation=_doc_interpolation, bound=_doc_bound)
+grid_push.__doc__ = grid_push.__doc__.format(
+    interpolation=_doc_interpolation, bound=_doc_bound)
+grid_count.__doc__ = grid_count.__doc__.format(
+    interpolation=_doc_interpolation, bound=_doc_bound)
+grid_grad.__doc__ = grid_grad.__doc__.format(
+    interpolation=_doc_interpolation, bound=_doc_bound)
+spline_coeff.__doc__ = spline_coeff.__doc__.format(
+    interpolation=_doc_interpolation, bound=_doc_bound_coeff, ref=_ref_coeff)
+spline_coeff_nd.__doc__ = spline_coeff_nd.__doc__.format(
+    interpolation=_doc_interpolation, bound=_doc_bound_coeff, ref=_ref_coeff)
+
+# aliases
+pull = grid_pull
+push = grid_push
+count = grid_count
+
+
+def identity_grid(shape, dtype=None, device=None):
+    """Returns an identity deformation field.
+
+    Parameters
+    ----------
+    shape : (dim,) sequence of int
+        Spatial dimension of the field.
+    dtype : torch.dtype, default=`get_default_dtype()`
+        Data type.
+    device torch.device, optional
+        Device.
+
+    Returns
+    -------
+    grid : (*shape, dim) tensor
+        Transformation field
+
+    """
+    mesh1d = [torch.arange(float(s), dtype=dtype, device=device)
+              for s in shape]
+    grid = torch.stack(meshgrid(mesh1d), dim=-1)
+    return grid
+
+
+@torch.jit.script
+def add_identity_grid_(disp):
+    """Adds the identity grid to a displacement field, inplace.
+
+    Parameters
+    ----------
+    disp : (..., *spatial, dim) tensor
+        Displacement field
+
+    Returns
+    -------
+    grid : (..., *spatial, dim) tensor
+        Transformation field
+
+    """
+    dim = disp.shape[-1]
+    spatial = disp.shape[-dim-1:-1]
+    mesh1d = [torch.arange(s, dtype=disp.dtype, device=disp.device)
+              for s in spatial]
+    grid = meshgrid(mesh1d)
+    disp = movedim1(disp, -1, 0)
+    for i, grid1 in enumerate(grid):
+        disp[i].add_(grid1)
+    disp = movedim1(disp, 0, -1)
+    return disp
+
+
+@torch.jit.script
+def add_identity_grid(disp):
+    """Adds the identity grid to a displacement field.
+
+    Parameters
+    ----------
+    disp : (..., *spatial, dim) tensor
+        Displacement field
+
+    Returns
+    -------
+    grid : (..., *spatial, dim) tensor
+        Transformation field
+
+    """
+    return add_identity_grid_(disp.clone())
+
+
+def affine_grid(mat, shape):
+    """Create a dense transformation grid from an affine matrix.
+
+    Parameters
+    ----------
+    mat : (..., D[+1], D+1) tensor
+        Affine matrix (or matrices).
+    shape : (D,) sequence[int]
+        Shape of the grid, with length D.
+
+    Returns
+    -------
+    grid : (..., *shape, D) tensor
+        Dense transformation grid
+
+    """
+    mat = torch.as_tensor(mat)
+    shape = list(shape)
+    nb_dim = mat.shape[-1] - 1
+    if nb_dim != len(shape):
+        raise ValueError('Dimension of the affine matrix ({}) and shape ({}) '
+                         'are not the same.'.format(nb_dim, len(shape)))
+    if mat.shape[-2] not in (nb_dim, nb_dim+1):
+        raise ValueError('First argument should be matrces of shape '
+                         '(..., {0}, {1}) or (..., {1], {1}) but got {2}.'
+                         .format(nb_dim, nb_dim+1, mat.shape))
+    batch_shape = mat.shape[:-2]
+    grid = identity_grid(shape, mat.dtype, mat.device)
+    if batch_shape:
+        for _ in range(len(batch_shape)):
+            grid = grid.unsqueeze(0)
+        for _ in range(nb_dim):
+            mat = mat.unsqueeze(-1)
+    lin = mat[..., :nb_dim, :nb_dim]
+    off = mat[..., :nb_dim, -1]
+    grid = matvec(lin, grid) + off
+    return grid

Generator/interpol/autograd.py

@@ -0,0 +1,301 @@
+"""AutoGrad version of pull/push/count/grad"""
+import torch
+from .coeff import spline_coeff_nd, spline_coeff
+from .bounds import BoundType
+from .splines import InterpolationType
+from .pushpull import (
+    grid_pull, grid_pull_backward,
+    grid_push, grid_push_backward,
+    grid_count, grid_count_backward,
+    grid_grad, grid_grad_backward)
+from .utils import fake_decorator
+try:
+    from torch.cuda.amp import custom_fwd, custom_bwd
+except (ModuleNotFoundError, ImportError):
+    custom_fwd = custom_bwd = fake_decorator
+
+
+def make_list(x):
+    if not isinstance(x, (list, tuple)):
+        x = [x]
+    return list(x)
+
+
+def bound_to_nitorch(bound, as_type='str'):
+    """Convert boundary type to niTorch's convention.
+
+    Parameters
+    ----------
+    bound : [list of] str or bound_like
+        Boundary condition in any convention
+    as_type : {'str', 'enum', 'int'}, default='str'
+        Return BoundType or int rather than str
+
+    Returns
+    -------
+    bound : [list of] str or BoundType
+        Boundary condition in NITorch's convention
+
+    """
+    intype = type(bound)
+    if not isinstance(bound, (list, tuple)):
+        bound = [bound]
+    obound = []
+    for b in bound:
+        b = b.lower() if isinstance(b, str) else b
+        if b in ('replicate', 'repeat', 'border', 'nearest', BoundType.replicate):
+            obound.append('replicate')
+        elif b in ('zero', 'zeros', 'constant', BoundType.zero):
+            obound.append('zero')
+        elif b in ('dct2', 'reflect', 'reflection', 'neumann', BoundType.dct2):
+            obound.append('dct2')
+        elif b in ('dct1', 'mirror', BoundType.dct1):
+            obound.append('dct1')
+        elif b in ('dft', 'wrap', 'circular', BoundType.dft):
+            obound.append('dft')
+        elif b in ('dst2', 'antireflect', 'dirichlet', BoundType.dst2):
+            obound.append('dst2')
+        elif b in ('dst1', 'antimirror', BoundType.dst1):
+            obound.append('dst1')
+        elif isinstance(b, int):
+            obound.append(b)
+        else:
+            raise ValueError(f'Unknown boundary condition {b}')
+    obound = list(map(lambda b: getattr(BoundType, b) if isinstance(b, str)
+                      else BoundType(b), obound))
+    if as_type in ('int', int):
+        obound = [b.value for b in obound]
+    if as_type in ('str', str):
+        obound = [b.name for b in obound]
+    if issubclass(intype, (list, tuple)):
+        obound = intype(obound)
+    else:
+        obound = obound[0]
+    return obound
+
+
+def inter_to_nitorch(inter, as_type='str'):
+    """Convert interpolation order to NITorch's convention.
+
+    Parameters
+    ----------
+    inter : [sequence of] int or str or InterpolationType
+    as_type : {'str', 'enum', 'int'}, default='int'
+
+    Returns
+    -------
+    inter : [sequence of] int or InterpolationType
+
+    """
+    intype = type(inter)
+    if not isinstance(inter, (list, tuple)):
+        inter = [inter]
+    ointer = []
+    for o in inter:
+        o = o.lower() if isinstance(o, str) else o
+        if o in (0, 'nearest', InterpolationType.nearest):
+            ointer.append(0)
+        elif o in (1, 'linear', InterpolationType.linear):
+            ointer.append(1)
+        elif o in (2, 'quadratic', InterpolationType.quadratic):
+            ointer.append(2)
+        elif o in (3, 'cubic', InterpolationType.cubic):
+            ointer.append(3)
+        elif o in (4, 'fourth', InterpolationType.fourth):
+            ointer.append(4)
+        elif o in (5, 'fifth', InterpolationType.fifth):
+            ointer.append(5)
+        elif o in (6, 'sixth', InterpolationType.sixth):
+            ointer.append(6)
+        elif o in (7, 'seventh', InterpolationType.seventh):
+            ointer.append(7)
+        else:
+            raise ValueError(f'Unknown interpolation order {o}')
+    if as_type in ('enum', 'str', str):
+        ointer = list(map(InterpolationType, ointer))
+        if as_type in ('str', str):
+            ointer = [o.name for o in ointer]
+    if issubclass(intype, (list, tuple)):
+        ointer = intype(ointer)
+    else:
+        ointer = ointer[0]
+    return ointer
+
+
+class GridPull(torch.autograd.Function):
+
+    @staticmethod
+    @custom_fwd(cast_inputs=torch.float32)
+    def forward(ctx, input, grid, interpolation, bound, extrapolate):
+
+        bound = bound_to_nitorch(make_list(bound), as_type='int')
+        interpolation = inter_to_nitorch(make_list(interpolation), as_type='int')
+        extrapolate = int(extrapolate)
+        opt = (bound, interpolation, extrapolate)
+
+        # Pull
+        output = grid_pull(input, grid, *opt)
+
+        # Context
+        ctx.opt = opt
+        ctx.save_for_backward(input, grid)
+
+        return output
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, grad):
+        var = ctx.saved_tensors
+        opt = ctx.opt
+        grads = grid_pull_backward(grad, *var, *opt)
+        grad_input, grad_grid = grads
+        return grad_input, grad_grid, None, None, None
+
+
+class GridPush(torch.autograd.Function):
+
+    @staticmethod
+    @custom_fwd(cast_inputs=torch.float32)
+    def forward(ctx, input, grid, shape, interpolation, bound, extrapolate):
+
+        bound = bound_to_nitorch(make_list(bound), as_type='int')
+        interpolation = inter_to_nitorch(make_list(interpolation), as_type='int')
+        extrapolate = int(extrapolate)
+        opt = (bound, interpolation, extrapolate)
+
+        # Push
+        output = grid_push(input, grid, shape, *opt)
+
+        # Context
+        ctx.opt = opt
+        ctx.save_for_backward(input, grid)
+
+        return output
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, grad):
+        var = ctx.saved_tensors
+        opt = ctx.opt
+        grads = grid_push_backward(grad, *var, *opt)
+        grad_input, grad_grid = grads
+        return grad_input, grad_grid, None, None, None, None
+
+
+class GridCount(torch.autograd.Function):
+
+    @staticmethod
+    @custom_fwd(cast_inputs=torch.float32)
+    def forward(ctx, grid, shape, interpolation, bound, extrapolate):
+
+        bound = bound_to_nitorch(make_list(bound), as_type='int')
+        interpolation = inter_to_nitorch(make_list(interpolation), as_type='int')
+        extrapolate = int(extrapolate)
+        opt = (bound, interpolation, extrapolate)
+
+        # Push
+        output = grid_count(grid, shape, *opt)
+
+        # Context
+        ctx.opt = opt
+        ctx.save_for_backward(grid)
+
+        return output
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, grad):
+        var = ctx.saved_tensors
+        opt = ctx.opt
+        grad_grid = None
+        if ctx.needs_input_grad[0]:
+            grad_grid = grid_count_backward(grad, *var, *opt)
+        return grad_grid, None, None, None, None
+
+
+class GridGrad(torch.autograd.Function):
+
+    @staticmethod
+    @custom_fwd(cast_inputs=torch.float32)
+    def forward(ctx, input, grid, interpolation, bound, extrapolate):
+
+        bound = bound_to_nitorch(make_list(bound), as_type='int')
+        interpolation = inter_to_nitorch(make_list(interpolation), as_type='int')
+        extrapolate = int(extrapolate)
+        opt = (bound, interpolation, extrapolate)
+
+        # Pull
+        output = grid_grad(input, grid, *opt)
+
+        # Context
+        ctx.opt = opt
+        ctx.save_for_backward(input, grid)
+
+        return output
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, grad):
+        var = ctx.saved_tensors
+        opt = ctx.opt
+        grad_input = grad_grid = None
+        if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]:
+            grads = grid_grad_backward(grad, *var, *opt)
+            grad_input, grad_grid = grads
+        return grad_input, grad_grid, None, None, None
+
+
+class SplineCoeff(torch.autograd.Function):
+
+    @staticmethod
+    @custom_fwd
+    def forward(ctx, input, bound, interpolation, dim, inplace):
+
+        bound = bound_to_nitorch(make_list(bound)[0], as_type='int')
+        interpolation = inter_to_nitorch(make_list(interpolation)[0], as_type='int')
+        opt = (bound, interpolation, dim, inplace)
+
+        # Pull
+        output = spline_coeff(input, *opt)
+
+        # Context
+        if input.requires_grad:
+            ctx.opt = opt
+
+        return output
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, grad):
+        # symmetric filter -> backward == forward
+        # (I don't know if I can write into grad, so inplace=False to be safe)
+        grad = spline_coeff(grad, *ctx.opt[:-1], inplace=False)
+        return [grad] + [None] * 4
+
+
+class SplineCoeffND(torch.autograd.Function):
+
+    @staticmethod
+    @custom_fwd
+    def forward(ctx, input, bound, interpolation, dim, inplace):
+
+        bound = bound_to_nitorch(make_list(bound), as_type='int')
+        interpolation = inter_to_nitorch(make_list(interpolation), as_type='int')
+        opt = (bound, interpolation, dim, inplace)
+
+        # Pull
+        output = spline_coeff_nd(input, *opt)
+
+        # Context
+        if input.requires_grad:
+            ctx.opt = opt
+
+        return output
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, grad):
+        # symmetric filter -> backward == forward
+        # (I don't know if I can write into grad, so inplace=False to be safe)
+        grad = spline_coeff_nd(grad, *ctx.opt[:-1], inplace=False)
+        return grad, None, None, None, None

Generator/interpol/backend.py

@@ -0,0 +1 @@
+jitfields = False  # Whether to use jitfields if available

Generator/interpol/bounds.py

@@ -0,0 +1,89 @@
+import torch
+from enum import Enum
+from typing import Optional
+from .jit_utils import floor_div
+Tensor = torch.Tensor
+
+
+class BoundType(Enum):
+    zero = zeros = 0
+    replicate = nearest = 1
+    dct1 = mirror = 2
+    dct2 = reflect = 3
+    dst1 = antimirror = 4
+    dst2 = antireflect = 5
+    dft = wrap = 6
+
+
+class ExtrapolateType(Enum):
+    no = 0     # threshold: (0, n-1)
+    yes = 1
+    hist = 2   # threshold: (-0.5, n-0.5)
+
+
+@torch.jit.script
+class Bound:
+
+    def __init__(self, bound_type: int = 3):
+        self.type = bound_type
+
+    def index(self, i, n: int):
+        if self.type in (0, 1):  # zero / replicate
+            return i.clamp(min=0, max=n-1)
+        elif self.type in (3, 5):  # dct2 / dst2
+            n2 = n * 2
+            i = torch.where(i < 0, (-i-1).remainder(n2).neg().add(n2 - 1),
+                            i.remainder(n2))
+            i = torch.where(i >= n, -i + (n2 - 1), i)
+            return i
+        elif self.type == 2:  # dct1
+            if n == 1:
+                return torch.zeros(i.shape, dtype=i.dtype, device=i.device)
+            else:
+                n2 = (n - 1) * 2
+                i = i.abs().remainder(n2)
+                i = torch.where(i >= n, -i + n2, i)
+                return i
+        elif self.type == 4:  # dst1
+            n2 = 2 * (n + 1)
+            first = torch.zeros([1], dtype=i.dtype, device=i.device)
+            last = torch.full([1], n - 1, dtype=i.dtype, device=i.device)
+            i = torch.where(i < 0, -i - 2, i)
+            i = i.remainder(n2)
+            i = torch.where(i > n, -i + (n2 - 2), i)
+            i = torch.where(i == -1, first, i)
+            i = torch.where(i == n, last, i)
+            return i
+        elif self.type == 6:  # dft
+            return i.remainder(n)
+        else:
+            return i
+
+    def transform(self, i, n: int) -> Optional[Tensor]:
+        if self.type == 4:  # dst1
+            if n == 1:
+                return None
+            one = torch.ones([1], dtype=torch.int8, device=i.device)
+            zero = torch.zeros([1], dtype=torch.int8, device=i.device)
+            n2 = 2 * (n + 1)
+            i = torch.where(i < 0, -i + (n-1), i)
+            i = i.remainder(n2)
+            x = torch.where(i == 0, zero, one)
+            x = torch.where(i.remainder(n + 1) == n, zero, x)
+            i = floor_div(i, n+1)
+            x = torch.where(torch.remainder(i, 2) > 0, -x, x)
+            return x
+        elif self.type == 5:  # dst2
+            i = torch.where(i < 0, n - 1 - i, i)
+            x = torch.ones([1], dtype=torch.int8, device=i.device)
+            i = floor_div(i, n)
+            x = torch.where(torch.remainder(i, 2) > 0, -x, x)
+            return x
+        elif self.type == 0:  # zero
+            one = torch.ones([1], dtype=torch.int8, device=i.device)
+            zero = torch.zeros([1], dtype=torch.int8, device=i.device)
+            outbounds = ((i < 0) | (i >= n))
+            x = torch.where(outbounds, zero, one)
+            return x
+        else:
+            return None

Generator/interpol/coeff.py

@@ -0,0 +1,344 @@
+"""Compute spline interpolating coefficients
+
+These functions are ported from the C routines in SPM's bsplines.c
+by John Ashburner, which are themselves ports from Philippe Thevenaz's
+code. JA furthermore derived the initial conditions for the DFT ("wrap around")
+boundary conditions.
+
+Note that similar routines are available in scipy with boundary conditions
+DCT1 ("mirror"), DCT2 ("reflect") and DFT ("wrap"); all derived by P. Thevenaz,
+according to the comments. Our DCT2 boundary conditions are ported from
+scipy.
+
+Only boundary conditions DCT1, DCT2 and DFT are implemented.
+
+References
+----------
+..[1]  M. Unser, A. Aldroubi and M. Eden.
+       "B-Spline Signal Processing: Part I-Theory,"
+       IEEE Transactions on Signal Processing 41(2):821-832 (1993).
+..[2]  M. Unser, A. Aldroubi and M. Eden.
+       "B-Spline Signal Processing: Part II-Efficient Design and Applications,"
+       IEEE Transactions on Signal Processing 41(2):834-848 (1993).
+..[3]  M. Unser.
+       "Splines: A Perfect Fit for Signal and Image Processing,"
+       IEEE Signal Processing Magazine 16(6):22-38 (1999).
+"""
+import torch
+import math
+from typing import List, Optional
+from .jit_utils import movedim1
+from .pushpull import pad_list_int
+
+
+@torch.jit.script
+def get_poles(order: int) -> List[float]:
+    empty: List[float] = []
+    if order in (0, 1):
+        return empty
+    if order == 2:
+        return [math.sqrt(8.) - 3.]
+    if order == 3:
+        return [math.sqrt(3.) - 2.]
+    if order == 4:
+        return [math.sqrt(664. - math.sqrt(438976.)) + math.sqrt(304.) - 19.,
+                math.sqrt(664. + math.sqrt(438976.)) - math.sqrt(304.) - 19.]
+    if order == 5:
+        return [math.sqrt(67.5 - math.sqrt(4436.25)) + math.sqrt(26.25) - 6.5,
+                math.sqrt(67.5 + math.sqrt(4436.25)) - math.sqrt(26.25) - 6.5]
+    if order == 6:
+        return [-0.488294589303044755130118038883789062112279161239377608394,
+                -0.081679271076237512597937765737059080653379610398148178525368,
+                -0.00141415180832581775108724397655859252786416905534669851652709]
+    if order == 7:
+        return [-0.5352804307964381655424037816816460718339231523426924148812,
+                -0.122554615192326690515272264359357343605486549427295558490763,
+                -0.0091486948096082769285930216516478534156925639545994482648003]
+    raise NotImplementedError
+
+
+@torch.jit.script
+def get_gain(poles: List[float]) -> float:
+    lam: float = 1.
+    for pole in poles:
+        lam *= (1. - pole) * (1. - 1./pole)
+    return lam
+
+
+@torch.jit.script
+def dft_initial(inp, pole: float, dim: int = -1, keepdim: bool = False):
+
+    assert inp.shape[dim] > 1
+    max_iter: int = int(math.ceil(-30./math.log(abs(pole))))
+    max_iter = min(max_iter, inp.shape[dim])
+
+    poles = torch.as_tensor(pole, dtype=inp.dtype, device=inp.device)
+    poles = poles.pow(torch.arange(1, max_iter, dtype=inp.dtype, device=inp.device))
+    poles = poles.flip(0)
+
+    inp = movedim1(inp, dim, 0)
+    inp0 = inp[0]
+    inp = inp[1-max_iter:]
+    inp = movedim1(inp, 0, -1)
+    out = torch.matmul(inp.unsqueeze(-2), poles.unsqueeze(-1)).squeeze(-1)
+    out = out + inp0.unsqueeze(-1)
+    if keepdim:
+        out = movedim1(out, -1, dim)
+    else:
+        out = out.squeeze(-1)
+
+    pole = pole ** max_iter
+    out = out / (1 - pole)
+    return out
+
+
+@torch.jit.script
+def dct1_initial(inp, pole: float, dim: int = -1, keepdim: bool = False):
+
+    n = inp.shape[dim]
+    max_iter: int = int(math.ceil(-30./math.log(abs(pole))))
+
+    if max_iter < n:
+
+        poles = torch.as_tensor(pole, dtype=inp.dtype, device=inp.device)
+        poles = poles.pow(torch.arange(1, max_iter, dtype=inp.dtype, device=inp.device))
+
+        inp = movedim1(inp, dim, 0)
+        inp0 = inp[0]
+        inp = inp[1:max_iter]
+        inp = movedim1(inp, 0, -1)
+        out = torch.matmul(inp.unsqueeze(-2), poles.unsqueeze(-1)).squeeze(-1)
+        out = out + inp0.unsqueeze(-1)
+        if keepdim:
+            out = movedim1(out, -1, dim)
+        else:
+            out = out.squeeze(-1)
+
+    else:
+        max_iter = n
+
+        polen = pole ** (n - 1)
+        inp0 = inp[0] + polen * inp[-1]
+        inp = inp[1:-1]
+        inp = movedim1(inp, 0, -1)
+
+        poles = torch.as_tensor(pole, dtype=inp.dtype, device=inp.device)
+        poles = poles.pow(torch.arange(1, n-1, dtype=inp.dtype, device=inp.device))
+        poles = poles + (polen * polen) / poles
+
+        out = torch.matmul(inp.unsqueeze(-2), poles.unsqueeze(-1)).squeeze(-1)
+        out = out + inp0.unsqueeze(-1)
+        if keepdim:
+            out = movedim1(out, -1, dim)
+        else:
+            out = out.squeeze(-1)
+
+        pole = pole ** (max_iter - 1)
+        out = out / (1 - pole * pole)
+
+    return out
+
+
+@torch.jit.script
+def dct2_initial(inp, pole: float, dim: int = -1, keepdim: bool = False):
+    # Ported from scipy:
+    # https://github.com/scipy/scipy/blob/master/scipy/ndimage/src/ni_splines.c
+    #
+    # I (YB) unwarped and simplied the terms so that I could use a dot
+    # product instead of a loop.
+    # It should certainly be possible to derive a version for max_iter < n,
+    # as JA did for DCT1, to avoid long recursions when `n` is large. But
+    # I think it would require a more complicated anticausal/final condition.
+
+    n = inp.shape[dim]
+
+    polen = pole ** n
+    pole_last = polen * (1 + 1/(pole + polen * polen))
+    inp00 = inp[0]
+    inp0 = inp[0] + pole_last * inp[-1]
+    inp = inp[1:-1]
+    inp = movedim1(inp, 0, -1)
+
+    poles = torch.as_tensor(pole, dtype=inp.dtype, device=inp.device)
+    poles = (poles.pow(torch.arange(1, n-1, dtype=inp.dtype, device=inp.device)) +
+             poles.pow(torch.arange(2*n-2, n, -1, dtype=inp.dtype, device=inp.device)))
+
+    out = torch.matmul(inp.unsqueeze(-2), poles.unsqueeze(-1)).squeeze(-1)
+
+    out = out + inp0.unsqueeze(-1)
+    out = out * (pole / (1 - polen * polen))
+    out = out + inp00.unsqueeze(-1)
+
+    if keepdim:
+        out = movedim1(out, -1, dim)
+    else:
+        out = out.squeeze(-1)
+
+    return out
+
+
+@torch.jit.script
+def dft_final(inp, pole: float, dim: int = -1, keepdim: bool = False):
+
+    assert inp.shape[dim] > 1
+    max_iter: int = int(math.ceil(-30./math.log(abs(pole))))
+    max_iter = min(max_iter, inp.shape[dim])
+
+    poles = torch.as_tensor(pole, dtype=inp.dtype, device=inp.device)
+    poles = poles.pow(torch.arange(2, max_iter+1, dtype=inp.dtype, device=inp.device))
+
+    inp = movedim1(inp, dim, 0)
+    inp0 = inp[-1]
+    inp = inp[:max_iter-1]
+    inp = movedim1(inp, 0, -1)
+    out = torch.matmul(inp.unsqueeze(-2), poles.unsqueeze(-1)).squeeze(-1)
+    out = out.add(inp0.unsqueeze(-1), alpha=pole)
+    if keepdim:
+        out = movedim1(out, -1, dim)
+    else:
+        out = out.squeeze(-1)
+
+    pole = pole ** max_iter
+    out = out / (pole - 1)
+    return out
+
+
+@torch.jit.script
+def dct1_final(inp, pole: float, dim: int = -1, keepdim: bool = False):
+    inp = movedim1(inp, dim, 0)
+    out = pole * inp[-2] + inp[-1]
+    out = out * (pole / (pole*pole - 1))
+    if keepdim:
+        out = movedim1(out.unsqueeze(0), 0, dim)
+    return out
+
+
+@torch.jit.script
+def dct2_final(inp, pole: float, dim: int = -1, keepdim: bool = False):
+    # Ported from scipy:
+    # https://github.com/scipy/scipy/blob/master/scipy/ndimage/src/ni_splines.c
+    inp = movedim1(inp, dim, 0)
+    out = inp[-1] * (pole / (pole - 1))
+    if keepdim:
+        out = movedim1(out.unsqueeze(0), 0, dim)
+    return out
+
+
+@torch.jit.script
+class CoeffBound:
+
+    def __init__(self, bound: int):
+        self.bound = bound
+
+    def initial(self, inp, pole: float, dim: int = -1, keepdim: bool = False):
+        if self.bound in (0, 2):    # zero, dct1
+            return dct1_initial(inp, pole, dim, keepdim)
+        elif self.bound in (1, 3):  # nearest, dct2
+            return dct2_initial(inp, pole, dim, keepdim)
+        elif self.bound == 6:       # dft
+            return dft_initial(inp, pole, dim, keepdim)
+        else:
+            raise NotImplementedError
+
+    def final(self, inp, pole: float, dim: int = -1, keepdim: bool = False):
+        if self.bound in (0, 2):    # zero, dct1
+            return dct1_final(inp, pole, dim, keepdim)
+        elif self.bound in (1, 3):  # nearest, dct2
+            return dct2_final(inp, pole, dim, keepdim)
+        elif self.bound == 6:       # dft
+            return dft_final(inp, pole, dim, keepdim)
+        else:
+            raise NotImplementedError
+
+
+@torch.jit.script
+def filter(inp, bound: CoeffBound, poles: List[float],
+           dim: int = -1, inplace: bool = False):
+
+    if not inplace:
+        inp = inp.clone()
+
+    if inp.shape[dim] == 1:
+        return inp
+
+    gain = get_gain(poles)
+    inp *= gain
+    inp = movedim1(inp, dim, 0)
+    n = inp.shape[0]
+
+    for pole in poles:
+        inp[0] = bound.initial(inp, pole, dim=0, keepdim=False)
+
+        for i in range(1, n):
+            inp[i].add_(inp[i-1], alpha=pole)
+
+        inp[-1] = bound.final(inp, pole, dim=0, keepdim=False)
+
+        for i in range(n-2, -1, -1):
+            inp[i].neg_().add_(inp[i+1]).mul_(pole)
+
+    inp = movedim1(inp, 0, dim)
+    return inp
+
+
+@torch.jit.script
+def spline_coeff(inp, bound: int, order: int, dim: int = -1,
+                 inplace: bool = False):
+    """Compute the interpolating spline coefficients, for a given spline order
+    and boundary conditions, along a single dimension.
+
+    Parameters
+    ----------
+    inp : tensor
+    bound : {2: dct1, 6: dft}
+    order : {0..7}
+    dim : int, default=-1
+    inplace : bool, default=False
+
+    Returns
+    -------
+    out : tensor
+
+    """
+    if not inplace:
+        inp = inp.clone()
+
+    if order in (0, 1):
+        return inp
+
+    poles = get_poles(order)
+    return filter(inp, CoeffBound(bound), poles, dim=dim, inplace=True)
+
+
+@torch.jit.script
+def spline_coeff_nd(inp, bound: List[int], order: List[int],
+                    dim: Optional[int] = None, inplace: bool = False):
+    """Compute the interpolating spline coefficients, for a given spline order
+    and boundary condition, along the last `dim` dimensions.
+
+    Parameters
+    ----------
+    inp : (..., *spatial) tensor
+    bound : List[{2: dct1, 6: dft}]
+    order : List[{0..7}]
+    dim : int, default=`inp.dim()`
+    inplace : bool, default=False
+
+    Returns
+    -------
+    out : (..., *spatial) tensor
+
+    """
+    if not inplace:
+        inp = inp.clone()
+
+    if dim is None:
+        dim = inp.dim()
+
+    bound = pad_list_int(bound, dim)
+    order = pad_list_int(order, dim)
+
+    for d, b, o in zip(range(dim), bound, order):
+        inp = spline_coeff(inp, b, o, dim=-dim + d, inplace=True)
+
+    return inp

Generator/interpol/iso0.py

@@ -0,0 +1,368 @@
+"""Isotropic 0-th order splines ("nearest neighbor")"""
+import torch
+from .bounds import Bound
+from .jit_utils import (sub2ind_list, make_sign,
+                        inbounds_mask_3d, inbounds_mask_2d, inbounds_mask_1d)
+from typing import List, Optional
+Tensor = torch.Tensor
+
+
+@torch.jit.script
+def get_indices(g, n: int, bound: Bound):
+    g0 = g.round().long()
+    sign0 = bound.transform(g0, n)
+    g0 = bound.index(g0, n)
+    return g0, sign0
+
+
+# ======================================================================
+#                                 3D
+# ======================================================================
+
+
+@torch.jit.script
+def pull3d(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY, iZ) tensor
+    g: (B, oX, oY, oZ, 3) tensor
+    bound: List{3}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, oX, oY, oZ) tensor
+    """
+    dim = 3
+    boundx, boundy, boundz = bound
+    oshape = g.shape[-dim-1:-1]
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy, gz = g.unbind(-1)
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+    shape = inp.shape[-dim:]
+    nx, ny, nz = shape
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_3d(extrapolate, gx, gy, gz, nx, ny, nz)
+
+    # nearest integer coordinates
+    gx, signx = get_indices(gx, nx, boundx)
+    gy, signy = get_indices(gy, ny, boundy)
+    gz, signz = get_indices(gz, nz, boundz)
+
+    # gather
+    inp = inp.reshape(inp.shape[:2] + [-1])
+    idx = sub2ind_list([gx, gy, gz], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out = inp.gather(-1, idx)
+    sign = make_sign([signx, signy, signz])
+    if sign is not None:
+        out *= sign
+    if mask is not None:
+        out *= mask
+    out = out.reshape(out.shape[:2] + oshape)
+    return out
+
+
+@torch.jit.script
+def push3d(inp, g, shape: Optional[List[int]], bound: List[Bound],
+           extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY, iZ) tensor
+    g: (B, iX, iY, iZ, 3) tensor
+    shape: List{3}[int], optional
+    bound: List{3}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, *shape) tensor
+    """
+    dim = 3
+    boundx, boundy, boundz = bound
+    if inp.shape[-dim:] != g.shape[-dim-1:-1]:
+        raise ValueError('Input and grid should have the same spatial shape')
+    ishape = inp.shape[-dim:]
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy, gz = torch.unbind(g, -1)
+    inp = inp.reshape(inp.shape[:2] + [-1])
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+
+    if shape is None:
+        shape = ishape
+    nx, ny, nz = shape
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_3d(extrapolate, gx, gy, gz, nx, ny, nz)
+
+    # nearest integer coordinates
+    gx, signx = get_indices(gx, nx, boundx)
+    gy, signy = get_indices(gy, ny, boundy)
+    gz, signz = get_indices(gz, nz, boundz)
+
+    # scatter
+    out = torch.zeros([batch, channel, nx*ny*nz], dtype=inp.dtype, device=inp.device)
+    idx = sub2ind_list([gx, gy, gz], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    sign = make_sign([signx, signy, signz])
+    if sign is not None or mask is not None:
+        inp = inp.clone()
+    if sign is not None:
+        inp *= sign
+    if mask is not None:
+        inp *= mask
+    out.scatter_add_(-1, idx, inp)
+
+    out = out.reshape(out.shape[:2] + shape)
+    return out
+
+
+# ======================================================================
+#                                 2D
+# ======================================================================
+
+
+@torch.jit.script
+def pull2d(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY) tensor
+    g: (B, oX, oY, 2) tensor
+    bound: List{2}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, oX, oY) tensor
+    """
+    dim = 2
+    boundx, boundy = bound
+    oshape = g.shape[-dim-1:-1]
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy = g.unbind(-1)
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+    shape = inp.shape[-dim:]
+    nx, ny = shape
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_2d(extrapolate, gx, gy, nx, ny)
+
+    # nearest integer coordinates
+    gx, signx = get_indices(gx, nx, boundx)
+    gy, signy = get_indices(gy, ny, boundy)
+
+    # gather
+    inp = inp.reshape(inp.shape[:2] + [-1])
+    idx = sub2ind_list([gx, gy], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out = inp.gather(-1, idx)
+    sign = make_sign([signx, signy])
+    if sign is not None:
+        out = out * sign
+    if mask is not None:
+        out = mask * mask
+    out = out.reshape(out.shape[:2] + oshape)
+    return out
+
+
+@torch.jit.script
+def push2d(inp, g, shape: Optional[List[int]], bound: List[Bound],
+           extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY) tensor
+    g: (B, iX, iY, 2) tensor
+    shape: List{2}[int], optional
+    bound: List{2}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, *shape) tensor
+    """
+    dim = 2
+    boundx, boundy = bound
+    if inp.shape[-dim:] != g.shape[-dim-1:-1]:
+        raise ValueError('Input and grid should have the same spatial shape')
+    ishape = inp.shape[-dim:]
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy = torch.unbind(g, -1)
+    inp = inp.reshape(inp.shape[:2] + [-1])
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+
+    if shape is None:
+        shape = ishape
+    nx, ny = shape
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_2d(extrapolate, gx, gy, nx, ny)
+
+    # nearest integer coordinates
+    gx, signx = get_indices(gx, nx, boundx)
+    gy, signy = get_indices(gy, ny, boundy)
+
+    # scatter
+    out = torch.zeros([batch, channel, nx*ny], dtype=inp.dtype, device=inp.device)
+    idx = sub2ind_list([gx, gy], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    sign = make_sign([signx, signy])
+    if sign is not None or mask is not None:
+        inp = inp.clone()
+    if sign is not None:
+        inp = inp * sign
+    if mask is not None:
+        inp = inp * mask
+    out.scatter_add_(-1, idx, inp)
+
+    out = out.reshape(out.shape[:2] + shape)
+    return out
+
+
+# ======================================================================
+#                                 1D
+# ======================================================================
+
+
+@torch.jit.script
+def pull1d(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, iX) tensor
+    g: (B, oX, 1) tensor
+    bound: List{1}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, oX) tensor
+    """
+    dim = 1
+    boundx = bound[0]
+    oshape = g.shape[-dim-1:-1]
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx = g.squeeze(-1)
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+    shape = inp.shape[-dim:]
+    nx = shape[0]
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_1d(extrapolate, gx, nx)
+
+    # nearest integer coordinates
+    gx, signx = get_indices(gx, nx, boundx)
+
+    # gather
+    inp = inp.reshape(inp.shape[:2] + [-1])
+    idx = gx
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out = inp.gather(-1, idx)
+    sign = signx
+    if sign is not None:
+        out = out * sign
+    if mask is not None:
+        out = out * mask
+    out = out.reshape(out.shape[:2] + oshape)
+    return out
+
+
+@torch.jit.script
+def push1d(inp, g, shape: Optional[List[int]], bound: List[Bound],
+           extrapolate: int = 1):
+    """
+    inp: (B, C, iX) tensor
+    g: (B, iX, 1) tensor
+    shape: List{1}[int], optional
+    bound: List{1}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, *shape) tensor
+    """
+    dim = 1
+    boundx = bound[0]
+    if inp.shape[-dim:] != g.shape[-dim-1:-1]:
+        raise ValueError('Input and grid should have the same spatial shape')
+    ishape = inp.shape[-dim:]
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx = g.squeeze(-1)
+    inp = inp.reshape(inp.shape[:2] + [-1])
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+
+    if shape is None:
+        shape = ishape
+    nx = shape[0]
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_1d(extrapolate, gx, nx)
+
+    # nearest integer coordinates
+    gx, signx = get_indices(gx, nx, boundx)
+
+    # scatter
+    out = torch.zeros([batch, channel, nx], dtype=inp.dtype, device=inp.device)
+    idx = gx
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    sign = signx
+    if sign is not None or mask is not None:
+        inp = inp.clone()
+    if sign is not None:
+        inp = inp * sign
+    if mask is not None:
+        inp = inp * mask
+    out.scatter_add_(-1, idx, inp)
+
+    out = out.reshape(out.shape[:2] + shape)
+    return out
+
+
+# ======================================================================
+#                                 ND
+# ======================================================================
+
+
+@torch.jit.script
+def grad(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, *ishape) tensor
+    g: (B, *oshape, D) tensor
+    bound: List{D}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, *oshape, D) tensor
+    """
+    dim = g.shape[-1]
+    oshape = list(g.shape[-dim-1:-1])
+    batch = max(inp.shape[0], g.shape[0])
+    channel = inp.shape[1]
+
+    return torch.zeros([batch, channel] + oshape + [dim],
+                       dtype=inp.dtype, device=inp.device)
+
+
+@torch.jit.script
+def pushgrad(inp, g, shape: Optional[List[int]], bound: List[Bound],
+             extrapolate: int = 1):
+    """
+    inp: (B, C, *ishape, D) tensor
+    g: (B, *ishape, D) tensor
+    shape: List{D}[int], optional, optional
+    bound: List{D}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, *shape) tensor
+    """
+    dim = g.shape[-1]
+    if inp.shape[-dim-1:-1] != g.shape[-dim-1:-1]:
+        raise ValueError('Input and grid should have the same spatial shape')
+    ishape = inp.shape[-dim-1:-1]
+    batch = max(inp.shape[0], g.shape[0])
+    channel = inp.shape[1]
+
+    if shape is None:
+        shape = ishape
+    shape = list(shape)
+
+    return torch.zeros([batch, channel] + shape,
+                       dtype=inp.dtype, device=inp.device)
+
+
+@torch.jit.script
+def hess(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, *ishape) tensor
+    g: (B, *oshape, D) tensor
+    bound: List{D}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, *oshape, D, D) tensor
+    """
+    dim = g.shape[-1]
+    oshape = list(g.shape[-dim-1:-1])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    batch = max(inp.shape[0], g.shape[0])
+    channel = inp.shape[1]
+
+    return torch.zeros([batch, channel] + oshape + [dim, dim],
+                       dtype=inp.dtype, device=inp.device)

Generator/interpol/iso1.py

@@ -0,0 +1,1339 @@
+"""Isotropic 1-st order splines ("linear/bilinear/trilinear")"""
+import torch
+from .bounds import Bound
+from .jit_utils import (sub2ind_list, make_sign,
+                        inbounds_mask_3d, inbounds_mask_2d, inbounds_mask_1d)
+from typing import List, Tuple, Optional
+Tensor = torch.Tensor
+
+
+@torch.jit.script
+def get_weights_and_indices(g, n: int, bound: Bound) \
+        -> Tuple[Tensor, Tensor, Tensor, Optional[Tensor], Optional[Tensor]]:
+    g0 = g.floor().long()
+    g1 = g0 + 1
+    sign1 = bound.transform(g1, n)
+    sign0 = bound.transform(g0, n)
+    g1 = bound.index(g1, n)
+    g0 = bound.index(g0, n)
+    g = g - g.floor()
+    return g, g0, g1, sign0, sign1
+
+
+# ======================================================================
+#                                 3D
+# ======================================================================
+
+
+@torch.jit.script
+def pull3d(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY, iZ) tensor
+    g: (B, oX, oY, oZ, 3) tensor
+    bound: List{3}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, oX, oY, oZ) tensor
+    """
+    dim = 3
+    boundx, boundy, boundz = bound
+    oshape = list(g.shape[-dim-1:-1])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy, gz = g.unbind(-1)
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+    shape = list(inp.shape[-dim:])
+    nx, ny, nz = shape
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_3d(extrapolate, gx, gy, gz, nx, ny, nz)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+    gy, gy0, gy1, signy0, signy1 = get_weights_and_indices(gy, ny, boundy)
+    gz, gz0, gz1, signz0, signz1 = get_weights_and_indices(gz, nz, boundz)
+
+    # gather
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
+    # - corner 000
+    idx = sub2ind_list([gx0, gy0, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy0, signz0])
+    if sign is not None:
+        out = out * sign
+    out = out * ((1 - gx) * (1 - gy) * (1 - gz))
+    # - corner 001
+    idx = sub2ind_list([gx0, gy0, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy0, signz1])
+    if sign is not None:
+        out1 = out1 * sign
+    out1 = out1 * ((1 - gx) * (1 - gy) * gz)
+    out = out + out1
+    # - corner 010
+    idx = sub2ind_list([gx0, gy1, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy1, signz0])
+    if sign is not None:
+        out1 = out1 * sign
+    out1 = out1 * ((1 - gx) * gy * (1 - gz))
+    out = out + out1
+    # - corner 011
+    idx = sub2ind_list([gx0, gy1, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy1, signz1])
+    if sign is not None:
+        out1 = out1 * sign
+    out1 = out1 * ((1 - gx) * gy * gz)
+    out = out + out1
+    # - corner 100
+    idx = sub2ind_list([gx1, gy0, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy0, signz0])
+    if sign is not None:
+        out1 = out1 * sign
+    out1 = out1 * (gx * (1 - gy) * (1 - gz))
+    out = out + out1
+    # - corner 101
+    idx = sub2ind_list([gx1, gy0, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy0, signz1])
+    if sign is not None:
+        out1 = out1 * sign
+    out1 = out1 * (gx * (1 - gy) * gz)
+    out = out + out1
+    # - corner 110
+    idx = sub2ind_list([gx1, gy1, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy1, signz0])
+    if sign is not None:
+        out1 = out1 * sign
+    out1 = out1 * (gx * gy * (1 - gz))
+    out = out + out1
+    # - corner 111
+    idx = sub2ind_list([gx1, gy1, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy1, signz1])
+    if sign is not None:
+        out1 = out1 * sign
+    out1 = out1 * (gx * gy * gz)
+    out = out + out1
+
+    if mask is not None:
+        out *= mask
+    out = out.reshape(list(out.shape[:2]) + oshape)
+    return out
+
+
+@torch.jit.script
+def push3d(inp, g, shape: Optional[List[int]], bound: List[Bound],
+           extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY, iZ) tensor
+    g: (B, iX, iY, iZ, 3) tensor
+    shape: List{3}[int], optional
+    bound: List{3}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, *shape) tensor
+    """
+    dim = 3
+    boundx, boundy, boundz = bound
+    if inp.shape[-dim:] != g.shape[-dim-1:-1]:
+        raise ValueError('Input and grid should have the same spatial shape')
+    ishape = list(inp.shape[-dim:])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy, gz = torch.unbind(g, -1)
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+
+    if shape is None:
+        shape = ishape
+    shape = list(shape)
+    nx, ny, nz = shape
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_3d(extrapolate, gx, gy, gz, nx, ny, nz)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+    gy, gy0, gy1, signy0, signy1 = get_weights_and_indices(gy, ny, boundy)
+    gz, gz0, gz1, signz0, signz1 = get_weights_and_indices(gz, nz, boundz)
+
+    # scatter
+    out = torch.zeros([batch, channel, nx*ny*nz],
+                      dtype=inp.dtype, device=inp.device)
+    # - corner 000
+    idx = sub2ind_list([gx0, gy0, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx0, signy0, signz0])
+    if sign is not None:
+        out1 = out1 * sign
+    if mask is not None:
+        out1 = out1 * mask
+    out1 = out1 * ((1 - gx) * (1 - gy) * (1 - gz))
+    out.scatter_add_(-1, idx, out1)
+    # - corner 001
+    idx = sub2ind_list([gx0, gy0, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx0, signy0, signz1])
+    if sign is not None:
+        out1 = out1 * sign
+    if mask is not None:
+        out1 = out1 * mask
+    out1 = out1 * ((1 - gx) * (1 - gy) * gz)
+    out.scatter_add_(-1, idx, out1)
+    # - corner 010
+    idx = sub2ind_list([gx0, gy1, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx0, signy1, signz0])
+    if sign is not None:
+        out1 = out1 * sign
+    if mask is not None:
+        out1 = out1 * mask
+    out1 = out1 * ((1 - gx) * gy * (1 - gz))
+    out.scatter_add_(-1, idx, out1)
+    # - corner 011
+    idx = sub2ind_list([gx0, gy1, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx0, signy1, signz1])
+    if sign is not None:
+        out1 = out1 * sign
+    if mask is not None:
+        out1 = out1 * mask
+    out1 = out1 * ((1 - gx) * gy * gz)
+    out.scatter_add_(-1, idx, out1)
+    # - corner 100
+    idx = sub2ind_list([gx1, gy0, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx1, signy0, signz0])
+    if sign is not None:
+        out1 = out1 * sign
+    if mask is not None:
+        out1 = out1 * mask
+    out1 = out1 * (gx * (1 - gy) * (1 - gz))
+    out.scatter_add_(-1, idx, out1)
+    # - corner 101
+    idx = sub2ind_list([gx1, gy0, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx1, signy0, signz1])
+    if sign is not None:
+        out1 = out1 * sign
+    if mask is not None:
+        out1 = out1 * mask
+    out1 = out1 * (gx * (1 - gy) * gz)
+    out.scatter_add_(-1, idx, out1)
+    # - corner 110
+    idx = sub2ind_list([gx1, gy1, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx1, signy1, signz0])
+    if sign is not None:
+        out1 = out1 * sign
+    if mask is not None:
+        out1 = out1 * mask
+    out1 = out1 * (gx * gy * (1 - gz))
+    out.scatter_add_(-1, idx, out1)
+    # - corner 111
+    idx = sub2ind_list([gx1, gy1, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx1, signy1, signz1])
+    if sign is not None:
+        out1 = out1 * sign
+    if mask is not None:
+        out1 = out1 * mask
+    out1 = out1 * (gx * gy * gz)
+    out.scatter_add_(-1, idx, out1)
+
+    out = out.reshape(list(out.shape[:2]) + shape)
+    return out
+
+
+@torch.jit.script
+def grad3d(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY, iZ) tensor
+    g: (B, oX, oY, oZ, 3) tensor
+    bound: List{3}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, oX, oY, oZ, 3) tensor
+    """
+    dim = 3
+    boundx, boundy, boundz = bound
+    oshape = list(g.shape[-dim-1:-1])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy, gz = torch.unbind(g, -1)
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+    shape = list(inp.shape[-dim:])
+    nx, ny, nz = shape
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_3d(extrapolate, gx, gy, gz, nx, ny, nz)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+    gy, gy0, gy1, signy0, signy1 = get_weights_and_indices(gy, ny, boundy)
+    gz, gz0, gz1, signz0, signz1 = get_weights_and_indices(gz, nz, boundz)
+
+    # gather
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
+    out = torch.empty([batch, channel] + list(g.shape[-2:]),
+                      dtype=inp.dtype, device=inp.device)
+    outx, outy, outz = out.unbind(-1)
+    # - corner 000
+    idx = sub2ind_list([gx0, gy0, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    torch.gather(inp, -1, idx, out=outx)
+    outy.copy_(outx)
+    outz.copy_(outx)
+    sign = make_sign([signx0, signy0, signz0])
+    if sign is not None:
+        out *= sign.unsqueeze(-1)
+    outx *= - (1 - gy) * (1 - gz)
+    outy *= - (1 - gx) * (1 - gz)
+    outz *= - (1 - gx) * (1 - gy)
+    # - corner 001
+    idx = sub2ind_list([gx0, gy0, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy0, signz1])
+    if sign is not None:
+        out1 *= sign
+    outx.addcmul_(out1, - (1 - gy) * gz)
+    outy.addcmul_(out1, - (1 - gx) * gz)
+    outz.addcmul_(out1,   (1 - gx) * (1 - gy))
+    # - corner 010
+    idx = sub2ind_list([gx0, gy1, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy1, signz0])
+    if sign is not None:
+        out1 *= sign
+    outx.addcmul_(out1, - gy * (1 - gz))
+    outy.addcmul_(out1, (1 - gx) * (1 - gz))
+    outz.addcmul_(out1, - (1 - gx) * gy)
+    # - corner 011
+    idx = sub2ind_list([gx0, gy1, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy1, signz1])
+    if sign is not None:
+        out1 *= sign
+    outx.addcmul_(out1, - gy * gz)
+    outy.addcmul_(out1, (1 - gx) * gz)
+    outz.addcmul_(out1, (1 - gx) * gy)
+    # - corner 100
+    idx = sub2ind_list([gx1, gy0, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy0, signz0])
+    if sign is not None:
+        out1 *= sign
+    outx.addcmul_(out1, (1 - gy) * (1 - gz))
+    outy.addcmul_(out1, - gx * (1 - gz))
+    outz.addcmul_(out1, - gx * (1 - gy))
+    # - corner 101
+    idx = sub2ind_list([gx1, gy0, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy0, signz1])
+    if sign is not None:
+        out1 *= sign
+    outx.addcmul_(out1, (1 - gy) * gz)
+    outy.addcmul_(out1, - gx * gz)
+    outz.addcmul_(out1, gx * (1 - gy))
+    # - corner 110
+    idx = sub2ind_list([gx1, gy1, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy1, signz0])
+    if sign is not None:
+        out1 *= sign
+    outx.addcmul_(out1, gy * (1 - gz))
+    outy.addcmul_(out1, gx * (1 - gz))
+    outz.addcmul_(out1, - gx * gy)
+    # - corner 111
+    idx = sub2ind_list([gx1, gy1, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy1, signz1])
+    if sign is not None:
+        out1 *= sign
+    outx.addcmul_(out1, gy * gz)
+    outy.addcmul_(out1, gx * gz)
+    outz.addcmul_(out1, gx * gy)
+
+    if mask is not None:
+        out *= mask.unsqueeze(-1)
+    out = out.reshape(list(out.shape[:2]) + oshape + [3])
+    return out
+
+
+@torch.jit.script
+def pushgrad3d(inp, g, shape: Optional[List[int]], bound: List[Bound],
+               extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY, iZ, 3) tensor
+    g: (B, iX, iY, iZ, 3) tensor
+    shape: List{3}[int], optional
+    bound: List{3}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, *shape) tensor
+    """
+    dim = 3
+    boundx, boundy, boundz = bound
+    if inp.shape[-dim-1:-1] != g.shape[-dim-1:-1]:
+        raise ValueError('Input and grid should have the same spatial shape')
+    ishape = list(inp.shape[-dim-1:-1])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy, gz = g.unbind(-1)
+    inp = inp.reshape(list(inp.shape[:2]) + [-1, dim])
+    batch = max(inp.shape[0], g.shape[0])
+    channel = inp.shape[1]
+
+    if shape is None:
+        shape = ishape
+    shape = list(shape)
+    nx, ny, nz = shape
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_3d(extrapolate, gx, gy, gz, nx, ny, nz)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+    gy, gy0, gy1, signy0, signy1 = get_weights_and_indices(gy, ny, boundy)
+    gz, gz0, gz1, signz0, signz1 = get_weights_and_indices(gz, nz, boundz)
+
+    # scatter
+    out = torch.zeros([batch, channel, nx*ny*nz],
+                      dtype=inp.dtype, device=inp.device)
+    # - corner 000
+    idx = sub2ind_list([gx0, gy0, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx0, signy0, signz0])
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x, out1y, out1z = out1.unbind(-1)
+    out1x *= - (1 - gy) * (1 - gz)
+    out1y *= - (1 - gx) * (1 - gz)
+    out1z *= - (1 - gx) * (1 - gy)
+    out.scatter_add_(-1, idx, out1x + out1y + out1z)
+    # - corner 001
+    idx = sub2ind_list([gx0, gy0, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx0, signy0, signz1])
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x, out1y, out1z = out1.unbind(-1)
+    out1x *= - (1 - gy) * gz
+    out1y *= - (1 - gx) * gz
+    out1z *= (1 - gx) * (1 - gy)
+    out.scatter_add_(-1, idx, out1x + out1y + out1z)
+    # - corner 010
+    idx = sub2ind_list([gx0, gy1, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx0, signy1, signz0])
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x, out1y, out1z = out1.unbind(-1)
+    out1x *= - gy * (1 - gz)
+    out1y *= (1 - gx) * (1 - gz)
+    out1z *= - (1 - gx) * gy
+    out.scatter_add_(-1, idx, out1x + out1y + out1z)
+    # - corner 011
+    idx = sub2ind_list([gx0, gy1, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx0, signy1, signz1])
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x, out1y, out1z = out1.unbind(-1)
+    out1x *= - gy * gz
+    out1y *= (1 - gx) * gz
+    out1z *= (1 - gx) * gy
+    out.scatter_add_(-1, idx, out1x + out1y + out1z)
+    # - corner 100
+    idx = sub2ind_list([gx1, gy0, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx1, signy0, signz0])
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x, out1y, out1z = out1.unbind(-1)
+    out1x *= (1 - gy) * (1 - gz)
+    out1y *= - gx * (1 - gz)
+    out1z *= - gx * (1 - gy)
+    out.scatter_add_(-1, idx, out1x + out1y + out1z)
+    # - corner 101
+    idx = sub2ind_list([gx1, gy0, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx1, signy0, signz1])
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x, out1y, out1z = out1.unbind(-1)
+    out1x *= (1 - gy) * gz
+    out1y *= - gx * gz
+    out1z *= gx * (1 - gy)
+    out.scatter_add_(-1, idx, out1x + out1y + out1z)
+    # - corner 110
+    idx = sub2ind_list([gx1, gy1, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx1, signy1, signz0])
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x, out1y, out1z = out1.unbind(-1)
+    out1x *= gy * (1 - gz)
+    out1y *= gx * (1 - gz)
+    out1z *= - gx * gy
+    out.scatter_add_(-1, idx, out1x + out1y + out1z)
+    # - corner 111
+    idx = sub2ind_list([gx1, gy1, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx1, signy1, signz1])
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x, out1y, out1z = out1.unbind(-1)
+    out1x *= gy * gz
+    out1y *= gx * gz
+    out1z *= gx * gy
+    out.scatter_add_(-1, idx, out1x + out1y + out1z)
+
+    out = out.reshape(list(out.shape[:2]) + shape)
+    return out
+
+
+@torch.jit.script
+def hess3d(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY, iZ) tensor
+    g: (B, oX, oY, oZ, 3) tensor
+    bound: List{3}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, oX, oY, oZ, 3, 3) tensor
+    """
+    dim = 3
+    boundx, boundy, boundz = bound
+    oshape = list(g.shape[-dim-1:-1])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy, gz = torch.unbind(g, -1)
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+    shape = list(inp.shape[-dim:])
+    nx, ny, nz = shape
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_3d(extrapolate, gx, gy, gz, nx, ny, nz)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+    gy, gy0, gy1, signy0, signy1 = get_weights_and_indices(gy, ny, boundy)
+    gz, gz0, gz1, signz0, signz1 = get_weights_and_indices(gz, nz, boundz)
+
+    # gather
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
+    out = torch.empty([batch, channel, g.shape[-2], dim, dim],
+                      dtype=inp.dtype, device=inp.device)
+    outx, outy, outz = out.unbind(-1)
+    outxx, outyx, outzx = outx.unbind(-1)
+    outxy, outyy, outzy = outy.unbind(-1)
+    outxz, outyz, outzz = outz.unbind(-1)
+    # - corner 000
+    idx = sub2ind_list([gx0, gy0, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    torch.gather(inp, -1, idx, out=outxy)
+    outxz.copy_(outxy)
+    outyz.copy_(outxy)
+    outxx.zero_()
+    outyy.zero_()
+    outzz.zero_()
+    sign = make_sign([signx0, signy0, signz0])
+    if sign is not None:
+        out *= sign.unsqueeze(-1).unsqueeze(-1)
+    outxy *= (1 - gz)
+    outxz *= (1 - gy)
+    outyz *= (1 - gx)
+    # - corner 001
+    idx = sub2ind_list([gx0, gy0, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy0, signz1])
+    if sign is not None:
+        out1 *= sign
+    outxy.addcmul_(out1, gz)
+    outxz.addcmul_(out1, - (1 - gy))
+    outyz.addcmul_(out1, - (1 - gx))
+    # - corner 010
+    idx = sub2ind_list([gx0, gy1, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy1, signz0])
+    if sign is not None:
+        out1 *= sign
+    outxy.addcmul_(out1, - (1 - gz))
+    outxz.addcmul_(out1, gy)
+    outyz.addcmul_(out1, - (1 - gx))
+    # - corner 011
+    idx = sub2ind_list([gx0, gy1, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy1, signz1])
+    if sign is not None:
+        out1 *= sign
+    outxy.addcmul_(out1, - gz)
+    outxz.addcmul_(out1, - gy)
+    outyz.addcmul_(out1, (1 - gx))
+    # - corner 100
+    idx = sub2ind_list([gx1, gy0, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy0, signz0])
+    if sign is not None:
+        out1 *= sign
+    outxy.addcmul_(out1, - (1 - gz))
+    outxz.addcmul_(out1, - (1 - gy))
+    outyz.addcmul_(out1, gx)
+    # - corner 101
+    idx = sub2ind_list([gx1, gy0, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy0, signz1])
+    if sign is not None:
+        out1 *= sign
+    outxy.addcmul_(out1, - gz)
+    outxz.addcmul_(out1, (1 - gy))
+    outyz.addcmul_(out1, - gx)
+    # - corner 110
+    idx = sub2ind_list([gx1, gy1, gz0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy1, signz0])
+    if sign is not None:
+        out1 *= sign
+    outxy.addcmul_(out1, (1 - gz))
+    outxz.addcmul_(out1, - gy)
+    outyz.addcmul_(out1, - gx)
+    # - corner 111
+    idx = sub2ind_list([gx1, gy1, gz1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy1, signz1])
+    if sign is not None:
+        out1 *= sign
+    outxy.addcmul_(out1, gz)
+    outxz.addcmul_(out1, gy)
+    outyz.addcmul_(out1, gx)
+
+    outyx.copy_(outxy)
+    outzx.copy_(outxz)
+    outzy.copy_(outyz)
+
+    if mask is not None:
+        out *= mask.unsqueeze(-1).unsqueeze(-1)
+    out = out.reshape(list(out.shape[:2]) + oshape + [dim, dim])
+    return out
+
+
+# ======================================================================
+#                                 2D
+# ======================================================================
+
+
+@torch.jit.script
+def pull2d(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY) tensor
+    g: (B, oX, oY, 2) tensor
+    bound: List{2}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, oX, oY) tensor
+    """
+    dim = 2
+    boundx, boundy = bound
+    oshape = list(g.shape[-dim-1:-1])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy = g.unbind(-1)
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+    shape = list(inp.shape[-dim:])
+    nx, ny = shape
+    
+    # mask of inbounds voxels
+    mask = inbounds_mask_2d(extrapolate, gx, gy, nx, ny)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+    gy, gy0, gy1, signy0, signy1 = get_weights_and_indices(gy, ny, boundy)
+
+    # gather
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
+    # - corner 00
+    idx = sub2ind_list([gx0, gy0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy0])
+    if sign is not None:
+        out = out * sign
+    out = out * ((1 - gx) * (1 - gy))
+    # - corner 01
+    idx = sub2ind_list([gx0, gy1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy1])
+    if sign is not None:
+        out1 = out1 * sign
+    out1 = out1 * ((1 - gx) * gy)
+    out = out + out1
+    # - corner 10
+    idx = sub2ind_list([gx1, gy0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy0])
+    if sign is not None:
+        out1 = out1 * sign
+    out1 = out1 * (gx * (1 - gy))
+    out = out + out1
+    # - corner 11
+    idx = sub2ind_list([gx1, gy1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy1])
+    if sign is not None:
+        out1 = out1 * sign
+    out1 = out1 * (gx * gy)
+    out = out + out1
+
+    if mask is not None:
+        out *= mask
+    out = out.reshape(list(out.shape[:2]) + oshape)
+    return out
+
+
+@torch.jit.script
+def push2d(inp, g, shape: Optional[List[int]], bound: List[Bound],
+           extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY) tensor
+    g: (B, iX, iY, 2) tensor
+    shape: List{2}[int], optional
+    bound: List{2}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, *shape) tensor
+    """
+    dim = 2
+    boundx, boundy = bound
+    if inp.shape[-dim:] != g.shape[-dim-1:-1]:
+        raise ValueError('Input and grid should have the same spatial shape')
+    ishape = list(inp.shape[-dim:])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy = torch.unbind(g, -1)
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+
+    if shape is None:
+        shape = ishape
+    shape = list(shape)
+    nx, ny = shape
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_2d(extrapolate, gx, gy, nx, ny)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+    gy, gy0, gy1, signy0, signy1 = get_weights_and_indices(gy, ny, boundy)
+
+    # scatter
+    out = torch.zeros([batch, channel, nx*ny],
+                      dtype=inp.dtype, device=inp.device)
+    # - corner 00
+    idx = sub2ind_list([gx0, gy0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx0, signy0])
+    if sign is not None:
+        out1 *= sign
+    if mask is not None:
+        out1 *= mask
+    out1 *= (1 - gx) * (1 - gy)
+    out.scatter_add_(-1, idx, out1)
+    # - corner 01
+    idx = sub2ind_list([gx0, gy1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx0, signy1])
+    if sign is not None:
+        out1 *= sign
+    if mask is not None:
+        out1 *= mask
+    out1 *= (1 - gx) * gy
+    out.scatter_add_(-1, idx, out1)
+    # - corner 10
+    idx = sub2ind_list([gx1, gy0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx1, signy0])
+    if sign is not None:
+        out1 *= sign
+    if mask is not None:
+        out1 *= mask
+    out1 *= gx * (1 - gy)
+    out.scatter_add_(-1, idx, out1)
+    # - corner 11
+    idx = sub2ind_list([gx1, gy1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx1, signy1])
+    if sign is not None:
+        out1 *= sign
+    if mask is not None:
+        out1 *= mask
+    out1 *= gx * gy
+    out.scatter_add_(-1, idx, out1)
+
+    out = out.reshape(list(out.shape[:2]) + shape)
+    return out
+
+
+@torch.jit.script
+def grad2d(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY) tensor
+    g: (B, oX, oY, 2) tensor
+    bound: List{2}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, oX, oY, 2) tensor
+    """
+    dim = 2
+    boundx, boundy = bound
+    oshape = list(g.shape[-dim-1:-1])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy = torch.unbind(g, -1)
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+    shape = list(inp.shape[-dim:])
+    nx, ny = shape
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_2d(extrapolate, gx, gy, nx, ny)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+    gy, gy0, gy1, signy0, signy1 = get_weights_and_indices(gy, ny, boundy)
+
+    # gather
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
+    out = torch.empty([batch, channel] + list(g.shape[-2:]),
+                      dtype=inp.dtype, device=inp.device)
+    outx, outy = out.unbind(-1)
+    # - corner 00
+    idx = sub2ind_list([gx0, gy0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    torch.gather(inp, -1, idx, out=outx)
+    outy.copy_(outx)
+    sign = make_sign([signx0, signy0])
+    if sign is not None:
+        out *= sign.unsqueeze(-1)
+    outx *= - (1 - gy)
+    outy *= - (1 - gx)
+    # - corner 01
+    idx = sub2ind_list([gx0, gy1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy1])
+    if sign is not None:
+        out1 *= sign
+    outx.addcmul_(out1, - gy)
+    outy.addcmul_(out1, (1 - gx))
+    # - corner 10
+    idx = sub2ind_list([gx1, gy0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy0])
+    if sign is not None:
+        out1 *= sign
+    outx.addcmul_(out1, (1 - gy))
+    outy.addcmul_(out1, - gx)
+    # - corner 11
+    idx = sub2ind_list([gx1, gy1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy1])
+    if sign is not None:
+        out1 *= sign
+    outx.addcmul_(out1, gy)
+    outy.addcmul_(out1, gx)
+
+    if mask is not None:
+        out *= mask.unsqueeze(-1)
+    out = out.reshape(list(out.shape[:2]) + oshape + [dim])
+    return out
+
+
+@torch.jit.script
+def pushgrad2d(inp, g, shape: Optional[List[int]], bound: List[Bound],
+               extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY, 2) tensor
+    g: (B, iX, iY, 2) tensor
+    shape: List{2}[int], optional
+    bound: List{2}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, *shape) tensor
+    """
+    dim = 2
+    boundx, boundy = bound
+    if inp.shape[-dim-1:-1] != g.shape[-dim-1:-1]:
+        raise ValueError('Input and grid should have the same spatial shape')
+    ishape = list(inp.shape[-dim-1:-1])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy = g.unbind(-1)
+    inp = inp.reshape(list(inp.shape[:2]) + [-1, dim])
+    batch = max(inp.shape[0], g.shape[0])
+    channel = inp.shape[1]
+
+    if shape is None:
+        shape = ishape
+    shape = list(shape)
+    nx, ny = shape
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_2d(extrapolate, gx, gy, nx, ny)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+    gy, gy0, gy1, signy0, signy1 = get_weights_and_indices(gy, ny, boundy)
+
+    # scatter
+    out = torch.zeros([batch, channel, nx*ny],
+                      dtype=inp.dtype, device=inp.device)
+    # - corner 00
+    idx = sub2ind_list([gx0, gy0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx0, signy0])
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x, out1y = out1.unbind(-1)
+    out1x *= - (1 - gy)
+    out1y *= - (1 - gx)
+    out.scatter_add_(-1, idx, out1x + out1y)
+    # - corner 01
+    idx = sub2ind_list([gx0, gy1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx0, signy1])
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x, out1y = out1.unbind(-1)
+    out1x *= - gy
+    out1y *= (1 - gx)
+    out.scatter_add_(-1, idx, out1x + out1y)
+    # - corner 10
+    idx = sub2ind_list([gx1, gy0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx1, signy0])
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x, out1y = out1.unbind(-1)
+    out1x *= (1 - gy)
+    out1y *= - gx
+    out.scatter_add_(-1, idx, out1x + out1y)
+    # - corner 11
+    idx = sub2ind_list([gx1, gy1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = make_sign([signx1, signy1])
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x, out1y = out1.unbind(-1)
+    out1x *= gy
+    out1y *= gx
+    out.scatter_add_(-1, idx, out1x + out1y)
+
+    out = out.reshape(list(out.shape[:2]) + shape)
+    return out
+
+
+@torch.jit.script
+def hess2d(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY) tensor
+    g: (B, oX, oY, 2) tensor
+    bound: List{2}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, oX, oY, 2, 2) tensor
+    """
+    dim = 2
+    boundx, boundy = bound
+    oshape = list(g.shape[-dim-1:-1])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx, gy = torch.unbind(g, -1)
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+    shape = list(inp.shape[-dim:])
+    nx, ny = shape
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_2d(extrapolate, gx, gy, nx, ny)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+    gy, gy0, gy1, signy0, signy1 = get_weights_and_indices(gy, ny, boundy)
+
+    # gather
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
+    out = torch.empty([batch, channel, g.shape[-2], dim, dim],
+                      dtype=inp.dtype, device=inp.device)
+    outx, outy = out.unbind(-1)
+    outxx, outyx = outx.unbind(-1)
+    outxy, outyy = outy.unbind(-1)
+    # - corner 00
+    idx = sub2ind_list([gx0, gy0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    torch.gather(inp, -1, idx, out=outxy)
+    outxx.zero_()
+    outyy.zero_()
+    sign = make_sign([signx0, signy0])
+    if sign is not None:
+        out *= sign.unsqueeze(-1).unsqueeze(-1)
+    outxy *= 1
+    # - corner 01
+    idx = sub2ind_list([gx0, gy1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx0, signy1])
+    if sign is not None:
+        out1 *= sign
+    outxy.add_(out1, alpha=-1)
+    # - corner 10
+    idx = sub2ind_list([gx1, gy0], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy0])
+    if sign is not None:
+        out1 *= sign
+    outxy.add_(out1, alpha=-1)
+    # - corner 11
+    idx = sub2ind_list([gx1, gy1], shape)
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = make_sign([signx1, signy1])
+    if sign is not None:
+        out1 *= sign
+    outxy.add_(out1)
+
+    outyx.copy_(outxy)
+
+    if mask is not None:
+        out *= mask.unsqueeze(-1).unsqueeze(-1)
+    out = out.reshape(list(out.shape[:2]) + oshape + [dim, dim])
+    return out
+
+
+# ======================================================================
+#                                 1D
+# ======================================================================
+
+
+@torch.jit.script
+def pull1d(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, iX) tensor
+    g: (B, oX, 1) tensor
+    bound: List{1}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, oX) tensor
+    """
+    dim = 1
+    boundx = bound[0]
+    oshape = list(g.shape[-dim-1:-1])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx = g.squeeze(-1)
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+    shape = list(inp.shape[-dim:])
+    nx = shape[0]
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_1d(extrapolate, gx, nx)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+
+    # gather
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
+    # - corner 0
+    idx = gx0
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out = inp.gather(-1, idx)
+    sign = signx0
+    if sign is not None:
+        out = out * sign
+    out = out * (1 - gx)
+    # - corner 1
+    idx = gx1
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = signx1
+    if sign is not None:
+        out1 = out1 * sign
+    out1 = out1 * gx
+    out = out + out1
+
+    if mask is not None:
+        out *= mask
+    out = out.reshape(list(out.shape[:2]) + oshape)
+    return out
+
+
+@torch.jit.script
+def push1d(inp, g, shape: Optional[List[int]], bound: List[Bound],
+           extrapolate: int = 1):
+    """
+    inp: (B, C, iX, iY) tensor
+    g: (B, iX, iY, 2) tensor
+    shape: List{2}[int], optional
+    bound: List{2}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, *shape) tensor
+    """
+    dim = 1
+    boundx = bound[0]
+    if inp.shape[-dim:] != g.shape[-dim-1:-1]:
+        raise ValueError('Input and grid should have the same spatial shape')
+    ishape = list(inp.shape[-dim:])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx = g.squeeze(-1)
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+
+    if shape is None:
+        shape = ishape
+    shape = list(shape)
+    nx = shape[0]
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_1d(extrapolate, gx, nx)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+
+    # scatter
+    out = torch.zeros([batch, channel, nx],
+                      dtype=inp.dtype, device=inp.device)
+    # - corner 0
+    idx = gx0
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = signx0
+    if sign is not None:
+        out1 = out1 * sign
+    if mask is not None:
+        out1 = out1 * mask
+    out1 = out1 * (1 - gx)
+    out.scatter_add_(-1, idx, out1)
+    # - corner 1
+    idx = gx1
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = signx1
+    if sign is not None:
+        out1 = out1 * sign
+    if mask is not None:
+        out1 = out1 * mask
+    out1 = out1 * gx
+    out.scatter_add_(-1, idx, out1)
+
+    out = out.reshape(list(out.shape[:2]) + shape)
+    return out
+
+
+@torch.jit.script
+def grad1d(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, iX) tensor
+    g: (B, oX, 1) tensor
+    bound: List{1}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, oX, 1) tensor
+    """
+    dim = 1
+    boundx = bound[0]
+    oshape = list(g.shape[-dim-1:-1])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx = g.squeeze(-1)
+    batch = max(inp.shape[0], gx.shape[0])
+    channel = inp.shape[1]
+    shape = list(inp.shape[-dim:])
+    nx = shape[0]
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_1d(extrapolate, gx, nx)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+
+    # gather
+    inp = inp.reshape(list(inp.shape[:2]) + [-1])
+    out = torch.empty([batch, channel] + list(g.shape[-2:]),
+                      dtype=inp.dtype, device=inp.device)
+    outx = out.squeeze(-1)
+    # - corner 0
+    idx = gx0
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    torch.gather(inp, -1, idx, out=outx)
+    sign = signx0
+    if sign is not None:
+        out *= sign.unsqueeze(-1)
+    outx.neg_()
+    # - corner 1
+    idx = gx1
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.gather(-1, idx)
+    sign = signx1
+    if sign is not None:
+        out1 *= sign
+    outx.add_(out1)
+
+    if mask is not None:
+        out *= mask.unsqueeze(-1)
+    out = out.reshape(list(out.shape[:2]) + oshape + [dim])
+    return out
+
+
+@torch.jit.script
+def pushgrad1d(inp, g, shape: Optional[List[int]], bound: List[Bound],
+               extrapolate: int = 1):
+    """
+    inp: (B, C, iX, 1) tensor
+    g: (B, iX, 1) tensor
+    shape: List{1}[int], optional
+    bound: List{1}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, *shape) tensor
+    """
+    dim = 1
+    boundx = bound[0]
+    if inp.shape[-2] != g.shape[-2]:
+        raise ValueError('Input and grid should have the same spatial shape')
+    ishape = list(inp.shape[-dim-1:-1])
+    g = g.reshape([g.shape[0], 1, -1, dim])
+    gx = g.squeeze(-1)
+    inp = inp.reshape(list(inp.shape[:2]) + [-1, dim])
+    batch = max(inp.shape[0], g.shape[0])
+    channel = inp.shape[1]
+
+    if shape is None:
+        shape = ishape
+    shape = list(shape)
+    nx = shape[0]
+
+    # mask of inbounds voxels
+    mask = inbounds_mask_1d(extrapolate, gx, nx)
+
+    # corners
+    # (upper weight, lower corner, upper corner, lower sign, upper sign)
+    gx, gx0, gx1, signx0, signx1 = get_weights_and_indices(gx, nx, boundx)
+
+    # scatter
+    out = torch.zeros([batch, channel, nx], dtype=inp.dtype, device=inp.device)
+    # - corner 000
+    idx = gx0
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = signx0
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x = out1.squeeze(-1)
+    out1x.neg_()
+    out.scatter_add_(-1, idx, out1x)
+    # - corner 100
+    idx = gx1
+    idx = idx.expand([batch, channel, idx.shape[-1]])
+    out1 = inp.clone()
+    sign = signx1
+    if sign is not None:
+        out1 *= sign.unsqueeze(-1)
+    if mask is not None:
+        out1 *= mask.unsqueeze(-1)
+    out1x = out1.squeeze(-1)
+    out.scatter_add_(-1, idx, out1x)
+
+    out = out.reshape(list(out.shape[:2]) + shape)
+    return out
+
+
+@torch.jit.script
+def hess1d(inp, g, bound: List[Bound], extrapolate: int = 1):
+    """
+    inp: (B, C, iX) tensor
+    g: (B, oX, 1) tensor
+    bound: List{1}[Bound] tensor
+    extrapolate: ExtrapolateType
+    returns: (B, C, oX, 1, 1) tensor
+    """
+    batch = max(inp.shape[0], g.shape[0])
+    return torch.zeros([batch, inp.shape[1], g.shape[1], 1, 1],
+                       dtype=inp.dtype, device=inp.device)

Generator/interpol/jit_utils.py

@@ -0,0 +1,443 @@
+"""A lot of utility functions for TorchScript"""
+import torch
+import os
+from typing import List, Tuple, Optional
+from .utils import torch_version
+from torch import Tensor
+
+
+@torch.jit.script
+def pad_list_int(x: List[int], dim: int) -> List[int]:
+    if len(x) < dim:
+        x = x + x[-1:] * (dim - len(x))
+    if len(x) > dim:
+        x = x[:dim]
+    return x
+
+
+@torch.jit.script
+def pad_list_float(x: List[float], dim: int) -> List[float]:
+    if len(x) < dim:
+        x = x + x[-1:] * (dim - len(x))
+    if len(x) > dim:
+        x = x[:dim]
+    return x
+
+
+@torch.jit.script
+def pad_list_str(x: List[str], dim: int) -> List[str]:
+    if len(x) < dim:
+        x = x + x[-1:] * (dim - len(x))
+    if len(x) > dim:
+        x = x[:dim]
+    return x
+
+
+@torch.jit.script
+def list_any(x: List[bool]) -> bool:
+    for elem in x:
+        if elem:
+            return True
+    return False
+
+
+@torch.jit.script
+def list_all(x: List[bool]) -> bool:
+    for elem in x:
+        if not elem:
+            return False
+    return True
+
+
+@torch.jit.script
+def list_prod_int(x: List[int]) -> int:
+    if len(x) == 0:
+        return 1
+    x0 = x[0]
+    for x1 in x[1:]:
+        x0 = x0 * x1
+    return x0
+
+
+@torch.jit.script
+def list_sum_int(x: List[int]) -> int:
+    if len(x) == 0:
+        return 1
+    x0 = x[0]
+    for x1 in x[1:]:
+        x0 = x0 + x1
+    return x0
+
+
+@torch.jit.script
+def list_prod_tensor(x: List[Tensor]) -> Tensor:
+    if len(x) == 0:
+        empty: List[int] = []
+        return torch.ones(empty)
+    x0 = x[0]
+    for x1 in x[1:]:
+        x0 = x0 * x1
+    return x0
+
+
+@torch.jit.script
+def list_sum_tensor(x: List[Tensor]) -> Tensor:
+    if len(x) == 0:
+        empty: List[int] = []
+        return torch.ones(empty)
+    x0 = x[0]
+    for x1 in x[1:]:
+        x0 = x0 + x1
+    return x0
+
+
+@torch.jit.script
+def list_reverse_int(x: List[int]) -> List[int]:
+    if len(x) == 0:
+        return x
+    return [x[i] for i in range(-1, -len(x)-1, -1)]
+
+
+@torch.jit.script
+def list_cumprod_int(x: List[int], reverse: bool = False,
+                     exclusive: bool = False) -> List[int]:
+    if len(x) == 0:
+        lx: List[int] = []
+        return lx
+    if reverse:
+        x = list_reverse_int(x)
+
+    x0 = 1 if exclusive else x[0]
+    lx = [x0]
+    all_x = x[:-1] if exclusive else x[1:]
+    for x1 in all_x:
+        x0 = x0 * x1
+        lx.append(x0)
+    if reverse:
+        lx = list_reverse_int(lx)
+    return lx
+
+
+@torch.jit.script
+def movedim1(x, source: int, destination: int):
+    dim = x.dim()
+    source = dim + source if source < 0 else source
+    destination = dim + destination if destination < 0 else destination
+    permutation = [d for d in range(dim)]
+    permutation = permutation[:source] + permutation[source+1:]
+    permutation = permutation[:destination] + [source] + permutation[destination:]
+    return x.permute(permutation)
+
+
+@torch.jit.script
+def sub2ind(subs, shape: List[int]):
+    """Convert sub indices (i, j, k) into linear indices.
+
+    The rightmost dimension is the most rapidly changing one
+    -> if shape == [D, H, W], the strides are therefore [H*W, W, 1]
+
+    Parameters
+    ----------
+    subs : (D, ...) tensor
+        List of sub-indices. The first dimension is the number of dimension.
+        Each element should have the same number of elements and shape.
+    shape : (D,) list[int]
+        Size of each dimension. Its length should be the same as the
+        first dimension of ``subs``.
+
+    Returns
+    -------
+    ind : (...) tensor
+        Linear indices
+    """
+    subs = subs.unbind(0)
+    ind = subs[-1]
+    subs = subs[:-1]
+    ind = ind.clone()
+    stride = list_cumprod_int(shape[1:], reverse=True, exclusive=False)
+    for i, s in zip(subs, stride):
+        ind += i * s
+    return ind
+
+
+@torch.jit.script
+def sub2ind_list(subs: List[Tensor], shape: List[int]):
+    """Convert sub indices (i, j, k) into linear indices.
+
+    The rightmost dimension is the most rapidly changing one
+    -> if shape == [D, H, W], the strides are therefore [H*W, W, 1]
+
+    Parameters
+    ----------
+    subs : (D,) list[tensor]
+        List of sub-indices. The first dimension is the number of dimension.
+        Each element should have the same number of elements and shape.
+    shape : (D,) list[int]
+        Size of each dimension. Its length should be the same as the
+        first dimension of ``subs``.
+
+    Returns
+    -------
+    ind : (...) tensor
+        Linear indices
+    """
+    ind = subs[-1]
+    subs = subs[:-1]
+    ind = ind.clone()
+    stride = list_cumprod_int(shape[1:], reverse=True, exclusive=False)
+    for i, s in zip(subs, stride):
+        ind += i * s
+    return ind
+
+# floor_divide returns wrong results for negative values, because it truncates
+# instead of performing a proper floor. In recent version of pytorch, it is
+# advised to use div(..., rounding_mode='trunc'|'floor') instead.
+# Here, we only use floor_divide on positive values so we do not care.
+if torch_version('>=', [1, 8]):
+    @torch.jit.script
+    def floor_div(x, y) -> torch.Tensor:
+        return torch.div(x, y, rounding_mode='floor')
+    @torch.jit.script
+    def floor_div_int(x, y: int) -> torch.Tensor:
+        return torch.div(x, y, rounding_mode='floor')
+else:
+    @torch.jit.script
+    def floor_div(x, y) -> torch.Tensor:
+        return (x / y).floor_()
+    @torch.jit.script
+    def floor_div_int(x, y: int) -> torch.Tensor:
+        return (x / y).floor_()
+
+
+@torch.jit.script
+def ind2sub(ind, shape: List[int]):
+    """Convert linear indices into sub indices (i, j, k).
+
+    The rightmost dimension is the most rapidly changing one
+    -> if shape == [D, H, W], the strides are therefore [H*W, W, 1]
+
+    Parameters
+    ----------
+    ind : tensor_like
+        Linear indices
+    shape : (D,) vector_like
+        Size of each dimension.
+
+    Returns
+    -------
+    subs : (D, ...) tensor
+        Sub-indices.
+    """
+    stride = list_cumprod_int(shape, reverse=True, exclusive=True)
+    sub = ind.new_empty([len(shape)] + ind.shape)
+    sub.copy_(ind)
+    for d in range(len(shape)):
+        if d > 0:
+            sub[d] = torch.remainder(sub[d], stride[d-1])
+        sub[d] = floor_div_int(sub[d], stride[d])
+    return sub
+
+
+@torch.jit.script
+def inbounds_mask_3d(extrapolate: int, gx, gy, gz, nx: int, ny: int, nz: int) \
+        -> Optional[Tensor]:
+    # mask of inbounds voxels
+    mask: Optional[Tensor] = None
+    if extrapolate in (0, 2):  # no / hist
+        tiny = 5e-2
+        threshold = tiny
+        if extrapolate == 2:
+            threshold = 0.5 + tiny
+        mask = ((gx > -threshold) & (gx < nx - 1 + threshold) &
+                (gy > -threshold) & (gy < ny - 1 + threshold) &
+                (gz > -threshold) & (gz < nz - 1 + threshold))
+        return mask
+    return mask
+
+
+@torch.jit.script
+def inbounds_mask_2d(extrapolate: int, gx, gy, nx: int, ny: int) \
+        -> Optional[Tensor]:
+    # mask of inbounds voxels
+    mask: Optional[Tensor] = None
+    if extrapolate in (0, 2):  # no / hist
+        tiny = 5e-2
+        threshold = tiny
+        if extrapolate == 2:
+            threshold = 0.5 + tiny
+        mask = ((gx > -threshold) & (gx < nx - 1 + threshold) &
+                (gy > -threshold) & (gy < ny - 1 + threshold))
+        return mask
+    return mask
+
+
+@torch.jit.script
+def inbounds_mask_1d(extrapolate: int, gx, nx: int) -> Optional[Tensor]:
+    # mask of inbounds voxels
+    mask: Optional[Tensor] = None
+    if extrapolate in (0, 2):  # no / hist
+        tiny = 5e-2
+        threshold = tiny
+        if extrapolate == 2:
+            threshold = 0.5 + tiny
+        mask = (gx > -threshold) & (gx < nx - 1 + threshold)
+        return mask
+    return mask
+
+
+@torch.jit.script
+def make_sign(sign: List[Optional[Tensor]]) -> Optional[Tensor]:
+    is_none : List[bool] = [s is None for s in sign]
+    if list_all(is_none):
+        return None
+    filt_sign: List[Tensor] = []
+    for s in sign:
+        if s is not None:
+            filt_sign.append(s)
+    return list_prod_tensor(filt_sign)
+
+
+@torch.jit.script
+def square(x):
+    return x * x
+
+
+@torch.jit.script
+def square_(x):
+    return x.mul_(x)
+
+
+@torch.jit.script
+def cube(x):
+    return x * x * x
+
+
+@torch.jit.script
+def cube_(x):
+    return square_(x).mul_(x)
+
+
+@torch.jit.script
+def pow4(x):
+    return square(square(x))
+
+
+@torch.jit.script
+def pow4_(x):
+    return square_(square_(x))
+
+
+@torch.jit.script
+def pow5(x):
+    return x * pow4(x)
+
+
+@torch.jit.script
+def pow5_(x):
+    return pow4_(x).mul_(x)
+
+
+@torch.jit.script
+def pow6(x):
+    return square(cube(x))
+
+
+@torch.jit.script
+def pow6_(x):
+    return square_(cube_(x))
+
+
+@torch.jit.script
+def pow7(x):
+    return pow6(x) * x
+
+
+@torch.jit.script
+def pow7_(x):
+    return pow6_(x).mul_(x)
+
+
+@torch.jit.script
+def dot(x, y, dim: int = -1, keepdim: bool = False):
+    """(Batched) dot product along a dimension"""
+    x = movedim1(x, dim, -1).unsqueeze(-2)
+    y = movedim1(y, dim, -1).unsqueeze(-1)
+    d = torch.matmul(x, y).squeeze(-1).squeeze(-1)
+    if keepdim:
+        d.unsqueeze(dim)
+    return d
+
+
+@torch.jit.script
+def dot_multi(x, y, dim: List[int], keepdim: bool = False):
+    """(Batched) dot product along a dimension"""
+    for d in dim:
+        x = movedim1(x, d, -1)
+        y = movedim1(y, d, -1)
+    x = x.reshape(x.shape[:-len(dim)] + [1, -1])
+    y = y.reshape(x.shape[:-len(dim)] + [-1, 1])
+    dt = torch.matmul(x, y).squeeze(-1).squeeze(-1)
+    if keepdim:
+        for d in dim:
+            dt.unsqueeze(d)
+    return dt
+
+
+
+# cartesian_prod takes multiple inout tensors as input in eager mode
+# but takes a list of tensor in jit mode. This is a helper that works
+# in both cases.
+if not int(os.environ.get('PYTORCH_JIT', '1')):
+    cartesian_prod = lambda x: torch.cartesian_prod(*x)
+    if torch_version('>=', (1, 10)):
+        def meshgrid_ij(x: List[torch.Tensor]) -> List[torch.Tensor]:
+            return torch.meshgrid(*x, indexing='ij')
+        def meshgrid_xy(x: List[torch.Tensor]) -> List[torch.Tensor]:
+            return torch.meshgrid(*x, indexing='xy')
+    else:
+        def meshgrid_ij(x: List[torch.Tensor]) -> List[torch.Tensor]:
+            return torch.meshgrid(*x)
+        def meshgrid_xy(x: List[torch.Tensor]) -> List[torch.Tensor]:
+            grid = torch.meshgrid(*x)
+            if len(grid) > 1:
+                grid[0] = grid[0].transpose(0, 1)
+                grid[1] = grid[1].transpose(0, 1)
+            return grid
+
+else:
+    cartesian_prod = torch.cartesian_prod
+    if torch_version('>=', (1, 10)):
+        @torch.jit.script
+        def meshgrid_ij(x: List[torch.Tensor]) -> List[torch.Tensor]:
+            return torch.meshgrid(x, indexing='ij')
+        @torch.jit.script
+        def meshgrid_xy(x: List[torch.Tensor]) -> List[torch.Tensor]:
+            return torch.meshgrid(x, indexing='xy')
+    else:
+        @torch.jit.script
+        def meshgrid_ij(x: List[torch.Tensor]) -> List[torch.Tensor]:
+            return torch.meshgrid(x)
+        @torch.jit.script
+        def meshgrid_xyt(x: List[torch.Tensor]) -> List[torch.Tensor]:
+            grid = torch.meshgrid(x)
+            if len(grid) > 1:
+                grid[0] = grid[0].transpose(0, 1)
+                grid[1] = grid[1].transpose(0, 1)
+            return grid
+
+
+meshgrid = meshgrid_ij
+
+
+# In torch < 1.6, div applied to integer tensor performed a floor_divide
+# In torch > 1.6, it performs a true divide.
+# Floor division must be done using `floor_divide`, but it was buggy
+# until torch 1.13 (it was doing a trunc divide instead of a floor divide).
+# There was at some point a deprecation warning for floor_divide, but it
+# seems to have been lifted afterwards. In torch >= 1.13, floor_divide
+# performs a correct floor division.
+# Since we only apply floor_divide ot positive values, we are fine.
+if torch_version('<', (1, 6)):
+    floor_div = torch.div
+else:
+    floor_div = torch.floor_divide

Generator/interpol/jitfields.py

@@ -0,0 +1,95 @@
+try:
+    import jitfields
+    available = True
+except (ImportError, ModuleNotFoundError):
+    jitfields = None
+    available = False
+from .utils import make_list
+import torch
+
+
+def first2last(input, ndim):
+    insert = input.dim() <= ndim
+    if insert:
+        input = input.unsqueeze(-1)
+    else:
+        input = torch.movedim(input, -ndim-1, -1)
+    return input, insert
+
+
+def last2first(input, ndim, inserted, grad=False):
+    if inserted:
+        input = input.squeeze(-1 - grad)
+    else:
+        input = torch.movedim(input, -1 - grad, -ndim-1 - grad)
+    return input
+
+
+def grid_pull(input, grid, interpolation='linear', bound='zero',
+              extrapolate=False, prefilter=False):
+    ndim = grid.shape[-1]
+    input, inserted = first2last(input, ndim)
+    input = jitfields.pull(input, grid, order=interpolation, bound=bound,
+                           extrapolate=extrapolate, prefilter=prefilter)
+    input = last2first(input, ndim, inserted)
+    return input
+
+
+def grid_push(input, grid, shape=None, interpolation='linear', bound='zero',
+              extrapolate=False, prefilter=False):
+    ndim = grid.shape[-1]
+    input, inserted = first2last(input, ndim)
+    input = jitfields.push(input, grid, shape, order=interpolation, bound=bound,
+                           extrapolate=extrapolate, prefilter=prefilter)
+    input = last2first(input, ndim, inserted)
+    return input
+
+
+def grid_count(grid, shape=None, interpolation='linear', bound='zero',
+               extrapolate=False):
+    return jitfields.count(grid, shape, order=interpolation, bound=bound,
+                           extrapolate=extrapolate)
+
+
+def grid_grad(input, grid, interpolation='linear', bound='zero',
+              extrapolate=False, prefilter=False):
+    ndim = grid.shape[-1]
+    input, inserted = first2last(input, ndim)
+    input = jitfields.grad(input, grid, order=interpolation, bound=bound,
+                           extrapolate=extrapolate, prefilter=prefilter)
+    input = last2first(input, ndim, inserted, True)
+    return input
+
+
+def spline_coeff(input, interpolation='linear', bound='dct2', dim=-1,
+                 inplace=False):
+    func = jitfields.spline_coeff_ if inplace else jitfields.spline_coeff
+    return func(input, interpolation, bound=bound, dim=dim)
+
+
+def spline_coeff_nd(input, interpolation='linear', bound='dct2', dim=None,
+                    inplace=False):
+    func = jitfields.spline_coeff_nd_ if inplace else jitfields.spline_coeff_nd
+    return func(input, interpolation, bound=bound, ndim=dim)
+
+
+def resize(image, factor=None, shape=None, anchor='c',
+           interpolation=1, prefilter=True, **kwargs):
+    kwargs.setdefault('bound', 'nearest')
+    ndim = max(len(make_list(factor or [])),
+               len(make_list(shape or [])),
+               len(make_list(anchor or []))) or (image.dim() - 2)
+    return jitfields.resize(image, factor=factor, shape=shape, ndim=ndim,
+                            anchor=anchor, order=interpolation,
+                            bound=kwargs['bound'], prefilter=prefilter)
+
+
+def restrict(image, factor=None, shape=None, anchor='c',
+             interpolation=1, reduce_sum=False, **kwargs):
+    kwargs.setdefault('bound', 'nearest')
+    ndim = max(len(make_list(factor or [])),
+               len(make_list(shape or [])),
+               len(make_list(anchor or []))) or (image.dim() - 2)
+    return jitfields.restrict(image, factor=factor, shape=shape, ndim=ndim,
+                              anchor=anchor, order=interpolation,
+                              bound=kwargs['bound'], reduce_sum=reduce_sum)

Generator/interpol/nd.py

@@ -0,0 +1,464 @@
+"""Generic N-dimensional version: any combination of spline orders"""
+import torch
+from typing import List, Optional, Tuple
+from .bounds import Bound
+from .splines import Spline
+from .jit_utils import sub2ind_list, make_sign, list_prod_int, cartesian_prod
+Tensor = torch.Tensor
+
+
+@torch.jit.script
+def inbounds_mask(extrapolate: int, grid, shape: List[int])\
+        -> Optional[Tensor]:
+    # mask of inbounds voxels
+    mask: Optional[Tensor] = None
+    if extrapolate in (0, 2):  # no / hist
+        grid = grid.unsqueeze(1)
+        tiny = 5e-2
+        threshold = tiny
+        if extrapolate == 2:
+            threshold = 0.5 + tiny
+        mask = torch.ones(grid.shape[:-1],
+                          dtype=torch.bool, device=grid.device)
+        for grid1, shape1 in zip(grid.unbind(-1), shape):
+            mask = mask & (grid1 > -threshold)
+            mask = mask & (grid1 < shape1 - 1 + threshold)
+        return mask
+    return mask
+
+
+@torch.jit.script
+def get_weights(grid, bound: List[Bound], spline: List[Spline],
+                shape: List[int], grad: bool = False, hess: bool = False) \
+        -> Tuple[List[List[Tensor]],
+                 List[List[Optional[Tensor]]],
+                 List[List[Optional[Tensor]]],
+                 List[List[Tensor]],
+                 List[List[Optional[Tensor]]]]:
+
+    weights: List[List[Tensor]] = []
+    grads: List[List[Optional[Tensor]]] = []
+    hesss: List[List[Optional[Tensor]]] = []
+    coords: List[List[Tensor]] = []
+    signs: List[List[Optional[Tensor]]] = []
+    for g, b, s, n in zip(grid.unbind(-1), bound, spline, shape):
+        grid0 = (g - (s.order-1)/2).floor()
+        dist0 = g - grid0
+        grid0 = grid0.long()
+        nb_nodes = s.order + 1
+        subweights: List[Tensor] = []
+        subcoords: List[Tensor] = []
+        subgrads: List[Optional[Tensor]] = []
+        subhesss: List[Optional[Tensor]] = []
+        subsigns: List[Optional[Tensor]] = []
+        for node in range(nb_nodes):
+            grid1 = grid0 + node
+            sign1: Optional[Tensor] = b.transform(grid1, n)
+            subsigns.append(sign1)
+            grid1 = b.index(grid1, n)
+            subcoords.append(grid1)
+            dist1 = dist0 - node
+            weight1 = s.fastweight(dist1)
+            subweights.append(weight1)
+            grad1: Optional[Tensor] = None
+            if grad:
+                grad1 = s.fastgrad(dist1)
+            subgrads.append(grad1)
+            hess1: Optional[Tensor] = None
+            if hess:
+                hess1 = s.fasthess(dist1)
+            subhesss.append(hess1)
+        weights.append(subweights)
+        coords.append(subcoords)
+        signs.append(subsigns)
+        grads.append(subgrads)
+        hesss.append(subhesss)
+
+    return weights, grads, hesss, coords, signs
+
+
+@torch.jit.script
+def pull(inp, grid, bound: List[Bound], spline: List[Spline],
+         extrapolate: int = 1):
+    """
+    inp: (B, C, *ishape) tensor
+    g: (B, *oshape, D) tensor
+    bound: List{D}[Bound] tensor
+    spline: List{D}[Spline] tensor
+    extrapolate: int
+    returns: (B, C, *oshape) tensor
+    """
+
+    dim = grid.shape[-1]
+    shape = list(inp.shape[-dim:])
+    oshape = list(grid.shape[-dim-1:-1])
+    batch = max(inp.shape[0], grid.shape[0])
+    channel = inp.shape[1]
+
+    grid = grid.reshape([grid.shape[0], -1, grid.shape[-1]])
+    inp = inp.reshape([inp.shape[0], inp.shape[1], -1])
+    mask = inbounds_mask(extrapolate, grid, shape)
+
+    # precompute weights along each dimension
+    weights, _, _, coords, signs = get_weights(grid, bound, spline, shape, False, False)
+
+    # initialize
+    out = torch.zeros([batch, channel, grid.shape[1]],
+                      dtype=inp.dtype, device=inp.device)
+
+    # iterate across nodes/corners
+    range_nodes = [torch.as_tensor([d for d in range(n)])
+                   for n in [s.order + 1 for s in spline]]
+    if dim == 1:
+        # cartesian_prod does not work as expected when only one
+        # element is provided
+        all_nodes = range_nodes[0].unsqueeze(-1)
+    else:
+        all_nodes = cartesian_prod(range_nodes)
+    for nodes in all_nodes:
+        # gather
+        idx = [c[n] for c, n in zip(coords, nodes)]
+        idx = sub2ind_list(idx, shape).unsqueeze(1)
+        idx = idx.expand([batch, channel, idx.shape[-1]])
+        out1 = inp.gather(-1, idx)
+
+        # apply sign
+        sign0: List[Optional[Tensor]] = [sgn[n] for sgn, n in zip(signs, nodes)]
+        sign1: Optional[Tensor] = make_sign(sign0)
+        if sign1 is not None:
+            out1 = out1 * sign1.unsqueeze(1)
+
+        # apply weights
+        for weight, n in zip(weights, nodes):
+            out1 = out1 * weight[n].unsqueeze(1)
+
+        # accumulate
+        out = out + out1
+
+    # out-of-bounds mask
+    if mask is not None:
+        out = out * mask
+
+    out = out.reshape(list(out.shape[:2]) + oshape)
+    return out
+
+
+@torch.jit.script
+def push(inp, grid, shape: Optional[List[int]], bound: List[Bound],
+         spline: List[Spline], extrapolate: int = 1):
+    """
+    inp: (B, C, *ishape) tensor
+    g: (B, *ishape, D) tensor
+    shape: List{D}[int], optional
+    bound: List{D}[Bound] tensor
+    spline: List{D}[Spline] tensor
+    extrapolate: int
+    returns: (B, C, *oshape) tensor
+    """
+
+    dim = grid.shape[-1]
+    ishape = list(grid.shape[-dim - 1:-1])
+    if shape is None:
+        shape = ishape
+    shape = list(shape)
+    batch = max(inp.shape[0], grid.shape[0])
+    channel = inp.shape[1]
+
+    grid = grid.reshape([grid.shape[0], -1, grid.shape[-1]])
+    inp = inp.reshape([inp.shape[0], inp.shape[1], -1])
+    mask = inbounds_mask(extrapolate, grid, shape)
+
+    # precompute weights along each dimension
+    weights, _, _, coords, signs = get_weights(grid, bound, spline, shape)
+
+    # initialize
+    out = torch.zeros([batch, channel, list_prod_int(shape)],
+                      dtype=inp.dtype, device=inp.device)
+
+    # iterate across nodes/corners
+    range_nodes = [torch.as_tensor([d for d in range(n)])
+                   for n in [s.order + 1 for s in spline]]
+    if dim == 1:
+        # cartesian_prod does not work as expected when only one
+        # element is provided
+        all_nodes = range_nodes[0].unsqueeze(-1)
+    else:
+        all_nodes = cartesian_prod(range_nodes)
+    for nodes in all_nodes:
+
+        # gather
+        idx = [c[n] for c, n in zip(coords, nodes)]
+        idx = sub2ind_list(idx, shape).unsqueeze(1)
+        idx = idx.expand([batch, channel, idx.shape[-1]])
+        out1 = inp.clone()
+
+        # apply sign
+        sign0: List[Optional[Tensor]] = [sgn[n] for sgn, n in zip(signs, nodes)]
+        sign1: Optional[Tensor] = make_sign(sign0)
+        if sign1 is not None:
+            out1 = out1 * sign1.unsqueeze(1)
+
+        # out-of-bounds mask
+        if mask is not None:
+            out1 = out1 * mask
+
+        # apply weights
+        for weight, n in zip(weights, nodes):
+            out1 = out1 * weight[n].unsqueeze(1)
+
+        # accumulate
+        out.scatter_add_(-1, idx, out1)
+
+    out = out.reshape(list(out.shape[:2]) + shape)
+    return out
+
+
+@torch.jit.script
+def grad(inp, grid, bound: List[Bound], spline: List[Spline],
+         extrapolate: int = 1):
+    """
+    inp: (B, C, *ishape) tensor
+    grid: (B, *oshape, D) tensor
+    bound: List{D}[Bound] tensor
+    spline: List{D}[Spline] tensor
+    extrapolate: int
+    returns: (B, C, *oshape, D) tensor
+    """
+
+    dim = grid.shape[-1]
+    shape = list(inp.shape[-dim:])
+    oshape = list(grid.shape[-dim-1:-1])
+    batch = max(inp.shape[0], grid.shape[0])
+    channel = inp.shape[1]
+
+    grid = grid.reshape([grid.shape[0], -1, grid.shape[-1]])
+    inp = inp.reshape([inp.shape[0], inp.shape[1], -1])
+    mask = inbounds_mask(extrapolate, grid, shape)
+
+    # precompute weights along each dimension
+    weights, grads, _, coords, signs = get_weights(grid, bound, spline, shape,
+                                                   grad=True)
+
+    # initialize
+    out = torch.zeros([batch, channel, grid.shape[1], dim],
+                      dtype=inp.dtype, device=inp.device)
+
+    # iterate across nodes/corners
+    range_nodes = [torch.as_tensor([d for d in range(n)])
+                   for n in [s.order + 1 for s in spline]]
+    if dim == 1:
+        # cartesian_prod does not work as expected when only one
+        # element is provided
+        all_nodes = range_nodes[0].unsqueeze(-1)
+    else:
+        all_nodes = cartesian_prod(range_nodes)
+    for nodes in all_nodes:
+
+        # gather
+        idx = [c[n] for c, n in zip(coords, nodes)]
+        idx = sub2ind_list(idx, shape).unsqueeze(1)
+        idx = idx.expand([batch, channel, idx.shape[-1]])
+        out0 = inp.gather(-1, idx)
+
+        # apply sign
+        sign0: List[Optional[Tensor]] = [sgn[n] for sgn, n in zip(signs, nodes)]
+        sign1: Optional[Tensor] = make_sign(sign0)
+        if sign1 is not None:
+            out0 = out0 * sign1.unsqueeze(1)
+
+        for d in range(dim):
+            out1 = out0.clone()
+            # apply weights
+            for dd, (weight, grad1, n) in enumerate(zip(weights, grads, nodes)):
+                if d == dd:
+                    grad11 = grad1[n]
+                    if grad11 is not None:
+                        out1 = out1 * grad11.unsqueeze(1)
+                else:
+                    out1 = out1 * weight[n].unsqueeze(1)
+
+            # accumulate
+            out.unbind(-1)[d].add_(out1)
+
+    # out-of-bounds mask
+    if mask is not None:
+        out = out * mask.unsqueeze(-1)
+
+    out = out.reshape(list(out.shape[:2]) + oshape + list(out.shape[-1:]))
+    return out
+
+
+@torch.jit.script
+def pushgrad(inp, grid, shape: Optional[List[int]], bound: List[Bound],
+             spline: List[Spline], extrapolate: int = 1):
+    """
+    inp: (B, C, *ishape, D) tensor
+    g: (B, *ishape, D) tensor
+    shape: List{D}[int], optional
+    bound: List{D}[Bound] tensor
+    spline: List{D}[Spline] tensor
+    extrapolate: int
+    returns: (B, C, *shape) tensor
+    """
+    dim = grid.shape[-1]
+    oshape = list(grid.shape[-dim-1:-1])
+    if shape is None:
+        shape = oshape
+    shape = list(shape)
+    batch = max(inp.shape[0], grid.shape[0])
+    channel = inp.shape[1]
+
+    grid = grid.reshape([grid.shape[0], -1, grid.shape[-1]])
+    inp = inp.reshape([inp.shape[0], inp.shape[1], -1, dim])
+    mask = inbounds_mask(extrapolate, grid, shape)
+
+    # precompute weights along each dimension
+    weights, grads, _, coords, signs = get_weights(grid, bound, spline, shape, grad=True)
+
+    # initialize
+    out = torch.zeros([batch, channel, list_prod_int(shape)],
+                      dtype=inp.dtype, device=inp.device)
+
+    # iterate across nodes/corners
+    range_nodes = [torch.as_tensor([d for d in range(n)])
+                   for n in [s.order + 1 for s in spline]]
+    if dim == 1:
+        # cartesian_prod does not work as expected when only one
+        # element is provided
+        all_nodes = range_nodes[0].unsqueeze(-1)
+    else:
+        all_nodes = cartesian_prod(range_nodes)
+    for nodes in all_nodes:
+
+        # gather
+        idx = [c[n] for c, n in zip(coords, nodes)]
+        idx = sub2ind_list(idx, shape).unsqueeze(1)
+        idx = idx.expand([batch, channel, idx.shape[-1]])
+        out0 = inp.clone()
+
+        # apply sign
+        sign0: List[Optional[Tensor]] = [sgn[n] for sgn, n in zip(signs, nodes)]
+        sign1: Optional[Tensor] = make_sign(sign0)
+        if sign1 is not None:
+            out0 = out0 * sign1.unsqueeze(1).unsqueeze(-1)
+
+        # out-of-bounds mask
+        if mask is not None:
+            out0 = out0 * mask.unsqueeze(-1)
+
+        for d in range(dim):
+            out1 = out0.unbind(-1)[d].clone()
+            # apply weights
+            for dd, (weight, grad1, n) in enumerate(zip(weights, grads, nodes)):
+                if d == dd:
+                    grad11 = grad1[n]
+                    if grad11 is not None:
+                        out1 = out1 * grad11.unsqueeze(1)
+                else:
+                    out1 = out1 * weight[n].unsqueeze(1)
+
+            # accumulate
+            out.scatter_add_(-1, idx, out1)
+
+    out = out.reshape(list(out.shape[:2]) + shape)
+    return out
+
+
+@torch.jit.script
+def hess(inp, grid, bound: List[Bound], spline: List[Spline],
+         extrapolate: int = 1):
+    """
+    inp: (B, C, *ishape) tensor
+    grid: (B, *oshape, D) tensor
+    bound: List{D}[Bound] tensor
+    spline: List{D}[Spline] tensor
+    extrapolate: int
+    returns: (B, C, *oshape, D, D) tensor
+    """
+
+    dim = grid.shape[-1]
+    shape = list(inp.shape[-dim:])
+    oshape = list(grid.shape[-dim-1:-1])
+    batch = max(inp.shape[0], grid.shape[0])
+    channel = inp.shape[1]
+
+    grid = grid.reshape([grid.shape[0], -1, grid.shape[-1]])
+    inp = inp.reshape([inp.shape[0], inp.shape[1], -1])
+    mask = inbounds_mask(extrapolate, grid, shape)
+
+    # precompute weights along each dimension
+    weights, grads, hesss, coords, signs \
+        = get_weights(grid, bound, spline, shape, grad=True, hess=True)
+
+    # initialize
+    out = torch.zeros([batch, channel, grid.shape[1], dim, dim],
+                      dtype=inp.dtype, device=inp.device)
+
+    # iterate across nodes/corners
+    range_nodes = [torch.as_tensor([d for d in range(n)])
+                   for n in [s.order + 1 for s in spline]]
+    if dim == 1:
+        # cartesian_prod does not work as expected when only one
+        # element is provided
+        all_nodes = range_nodes[0].unsqueeze(-1)
+    else:
+        all_nodes = cartesian_prod(range_nodes)
+    for nodes in all_nodes:
+
+        # gather
+        idx = [c[n] for c, n in zip(coords, nodes)]
+        idx = sub2ind_list(idx, shape).unsqueeze(1)
+        idx = idx.expand([batch, channel, idx.shape[-1]])
+        out0 = inp.gather(-1, idx)
+
+        # apply sign
+        sign0: List[Optional[Tensor]] = [sgn[n] for sgn, n in zip(signs, nodes)]
+        sign1: Optional[Tensor] = make_sign(sign0)
+        if sign1 is not None:
+            out0 = out0 * sign1.unsqueeze(1)
+
+        for d in range(dim):
+            # -- diagonal --
+            out1 = out0.clone()
+
+            # apply weights
+            for dd, (weight, hess1, n) \
+                    in enumerate(zip(weights, hesss, nodes)):
+                if d == dd:
+                    hess11 = hess1[n]
+                    if hess11 is not None:
+                        out1 = out1 * hess11.unsqueeze(1)
+                else:
+                    out1 = out1 * weight[n].unsqueeze(1)
+
+            # accumulate
+            out.unbind(-1)[d].unbind(-1)[d].add_(out1)
+
+            # -- off diagonal --
+            for d2 in range(d+1, dim):
+                out1 = out0.clone()
+
+                # apply weights
+                for dd, (weight, grad1, n) \
+                        in enumerate(zip(weights, grads, nodes)):
+                    if dd in (d, d2):
+                        grad11 = grad1[n]
+                        if grad11 is not None:
+                            out1 = out1 * grad11.unsqueeze(1)
+                    else:
+                        out1 = out1 * weight[n].unsqueeze(1)
+
+                # accumulate
+                out.unbind(-1)[d].unbind(-1)[d2].add_(out1)
+
+    # out-of-bounds mask
+    if mask is not None:
+        out = out * mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
+
+    # fill lower triangle
+    for d in range(dim):
+        for d2 in range(d+1, dim):
+            out.unbind(-1)[d2].unbind(-1)[d].copy_(out.unbind(-1)[d].unbind(-1)[d2])
+
+    out = out.reshape(list(out.shape[:2]) + oshape + list(out.shape[-2:]))
+    return out

Generator/interpol/pushpull.py

@@ -0,0 +1,325 @@
+"""
+Non-differentiable forward/backward components.
+These components are put together in `interpol.autograd` to generate
+differentiable functions.
+
+Note
+----
+.. I removed @torch.jit.script from these entry-points because compiling
+   all possible combinations of bound+interpolation made the first call
+   extremely slow.
+.. I am not using the dot/multi_dot helpers even though they should be
+   more efficient that "multiply and sum" because I haven't had the time
+   to test them. It would be worth doing it.
+"""
+import torch
+from typing import List, Optional, Tuple
+from .jit_utils import list_all, dot, dot_multi, pad_list_int
+from .bounds import Bound
+from .splines import Spline
+from . import iso0, iso1, nd
+Tensor = torch.Tensor
+
+
+@torch.jit.script
+def make_bound(bound: List[int]) -> List[Bound]:
+    return [Bound(b) for b in bound]
+
+
+@torch.jit.script
+def make_spline(spline: List[int]) -> List[Spline]:
+    return [Spline(s) for s in spline]
+
+
+# @torch.jit.script
+def grid_pull(inp, grid, bound: List[int], interpolation: List[int],
+              extrapolate: int):
+    """
+    inp: (B, C, *spatial_in) tensor
+    grid: (B, *spatial_out, D) tensor
+    bound: List{D}[int] tensor
+    interpolation: List{D}[int]
+    extrapolate: int
+    returns: (B, C, *spatial_out) tensor
+    """
+    dim = grid.shape[-1]
+    bound = pad_list_int(bound, dim)
+    interpolation = pad_list_int(interpolation, dim)
+    bound_fn = make_bound(bound)
+    is_iso1 = list_all([order == 1 for order in interpolation])
+    if is_iso1:
+        if dim == 3:
+            return iso1.pull3d(inp, grid, bound_fn, extrapolate)
+        elif dim == 2:
+            return iso1.pull2d(inp, grid, bound_fn, extrapolate)
+        elif dim == 1:
+            return iso1.pull1d(inp, grid, bound_fn, extrapolate)
+    is_iso0 = list_all([order == 0 for order in interpolation])
+    if is_iso0:
+        if dim == 3:
+            return iso0.pull3d(inp, grid, bound_fn, extrapolate)
+        elif dim == 2:
+            return iso0.pull2d(inp, grid, bound_fn, extrapolate)
+        elif dim == 1:
+            return iso0.pull1d(inp, grid, bound_fn, extrapolate)
+    spline_fn = make_spline(interpolation)
+    return nd.pull(inp, grid, bound_fn, spline_fn, extrapolate)
+
+
+# @torch.jit.script
+def grid_push(inp, grid, shape: Optional[List[int]], bound: List[int],
+              interpolation: List[int], extrapolate: int):
+    """
+    inp: (B, C, *spatial_in) tensor
+    grid: (B, *spatial_in, D) tensor
+    shape: List{D}[int] tensor, optional, default=spatial_in
+    bound: List{D}[int] tensor
+    interpolation: List{D}[int]
+    extrapolate: int
+    returns: (B, C, *shape) tensor
+    """
+    dim = grid.shape[-1]
+    bound = pad_list_int(bound, dim)
+    interpolation = pad_list_int(interpolation, dim)
+    bound_fn = make_bound(bound)
+    is_iso1 = list_all([order == 1 for order in interpolation])
+    if is_iso1:
+        if dim == 3:
+            return iso1.push3d(inp, grid, shape, bound_fn, extrapolate)
+        elif dim == 2:
+            return iso1.push2d(inp, grid, shape, bound_fn, extrapolate)
+        elif dim == 1:
+            return iso1.push1d(inp, grid, shape, bound_fn, extrapolate)
+    is_iso0 = list_all([order == 0 for order in interpolation])
+    if is_iso0:
+        if dim == 3:
+            return iso0.push3d(inp, grid, shape, bound_fn, extrapolate)
+        elif dim == 2:
+            return iso0.push2d(inp, grid, shape, bound_fn, extrapolate)
+        elif dim == 1:
+            return iso0.push1d(inp, grid, shape, bound_fn, extrapolate)
+    spline_fn = make_spline(interpolation)
+    return nd.push(inp, grid, shape, bound_fn, spline_fn, extrapolate)
+
+
+# @torch.jit.script
+def grid_count(grid, shape: Optional[List[int]], bound: List[int],
+              interpolation: List[int], extrapolate: int):
+    """
+    grid: (B, *spatial_in, D) tensor
+    shape: List{D}[int] tensor, optional, default=spatial_in
+    bound: List{D}[int] tensor
+    interpolation: List{D}[int]
+    extrapolate: int
+    returns: (B, 1, *shape) tensor
+    """
+    dim = grid.shape[-1]
+    bound = pad_list_int(bound, dim)
+    interpolation = pad_list_int(interpolation, dim)
+    bound_fn = make_bound(bound)
+    gshape = list(grid.shape[-dim-1:-1])
+    if shape is None:
+        shape = gshape
+    inp = torch.ones([], dtype=grid.dtype, device=grid.device)
+    inp = inp.expand([len(grid), 1] + gshape)
+    is_iso1 = list_all([order == 1 for order in interpolation])
+    if is_iso1:
+        if dim == 3:
+            return iso1.push3d(inp, grid, shape, bound_fn, extrapolate)
+        elif dim == 2:
+            return iso1.push2d(inp, grid, shape, bound_fn, extrapolate)
+        elif dim == 1:
+            return iso1.push1d(inp, grid, shape, bound_fn, extrapolate)
+    is_iso0 = list_all([order == 0 for order in interpolation])
+    if is_iso0:
+        if dim == 3:
+            return iso0.push3d(inp, grid, shape, bound_fn, extrapolate)
+        elif dim == 2:
+            return iso0.push2d(inp, grid, shape, bound_fn, extrapolate)
+        elif dim == 1:
+            return iso0.push1d(inp, grid, shape, bound_fn, extrapolate)
+    spline_fn = make_spline(interpolation)
+    return nd.push(inp, grid, shape, bound_fn, spline_fn, extrapolate)
+
+
+# @torch.jit.script
+def grid_grad(inp, grid, bound: List[int], interpolation: List[int],
+              extrapolate: int):
+    """
+    inp: (B, C, *spatial_in) tensor
+    grid: (B, *spatial_out, D) tensor
+    bound: List{D}[int] tensor
+    interpolation: List{D}[int]
+    extrapolate: int
+    returns: (B, C, *spatial_out, D) tensor
+    """
+    dim = grid.shape[-1]
+    bound = pad_list_int(bound, dim)
+    interpolation = pad_list_int(interpolation, dim)
+    bound_fn = make_bound(bound)
+    is_iso1 = list_all([order == 1 for order in interpolation])
+    if is_iso1:
+        if dim == 3:
+            return iso1.grad3d(inp, grid, bound_fn, extrapolate)
+        elif dim == 2:
+            return iso1.grad2d(inp, grid, bound_fn, extrapolate)
+        elif dim == 1:
+            return iso1.grad1d(inp, grid, bound_fn, extrapolate)
+    is_iso0 = list_all([order == 0 for order in interpolation])
+    if is_iso0:
+        return iso0.grad(inp, grid, bound_fn, extrapolate)
+    spline_fn = make_spline(interpolation)
+    return nd.grad(inp, grid, bound_fn, spline_fn, extrapolate)
+
+
+# @torch.jit.script
+def grid_pushgrad(inp, grid, shape: List[int], bound: List[int],
+                  interpolation: List[int], extrapolate: int):
+    """ /!\ Used only in backward pass of grid_grad
+    inp: (B, C, *spatial_in, D) tensor
+    grid: (B, *spatial_in, D) tensor
+    shape: List{D}[int], optional
+    bound: List{D}[int] tensor
+    interpolation: List{D}[int]
+    extrapolate: int
+    returns: (B, C, *shape) tensor
+    """
+    dim = grid.shape[-1]
+    bound = pad_list_int(bound, dim)
+    interpolation = pad_list_int(interpolation, dim)
+    bound_fn = make_bound(bound)
+    is_iso1 = list_all([order == 1 for order in interpolation])
+    if is_iso1:
+        if dim == 3:
+            return iso1.pushgrad3d(inp, grid, shape, bound_fn, extrapolate)
+        elif dim == 2:
+            return iso1.pushgrad2d(inp, grid, shape, bound_fn, extrapolate)
+        elif dim == 1:
+            return iso1.pushgrad1d(inp, grid, shape, bound_fn, extrapolate)
+    is_iso0 = list_all([order == 0 for order in interpolation])
+    if is_iso0:
+        return iso0.pushgrad(inp, grid, shape, bound_fn, extrapolate)
+    spline_fn = make_spline(interpolation)
+    return nd.pushgrad(inp, grid, shape, bound_fn, spline_fn, extrapolate)
+
+
+# @torch.jit.script
+def grid_hess(inp, grid, bound: List[int], interpolation: List[int],
+              extrapolate: int):
+    """ /!\ Used only in backward pass of grid_grad
+    inp: (B, C, *spatial_in) tensor
+    grid: (B, *spatial_out, D) tensor
+    bound: List{D}[int] tensor
+    interpolation: List{D}[int]
+    extrapolate: int
+    returns: (B, C, *spatial_out, D, D) tensor
+    """
+    dim = grid.shape[-1]
+    bound = pad_list_int(bound, dim)
+    interpolation = pad_list_int(interpolation, dim)
+    bound_fn = make_bound(bound)
+    is_iso1 = list_all([order == 1 for order in interpolation])
+    if is_iso1:
+        if dim == 3:
+            return iso1.hess3d(inp, grid, bound_fn, extrapolate)
+        if dim == 2:
+            return iso1.hess2d(inp, grid, bound_fn, extrapolate)
+        if dim == 1:
+            return iso1.hess1d(inp, grid, bound_fn, extrapolate)
+    is_iso0 = list_all([order == 0 for order in interpolation])
+    if is_iso0:
+        return iso0.hess(inp, grid, bound_fn, extrapolate)
+    spline_fn = make_spline(interpolation)
+    return nd.hess(inp, grid, bound_fn, spline_fn, extrapolate)
+
+
+# @torch.jit.script
+def grid_pull_backward(grad, inp, grid, bound: List[int],
+                       interpolation: List[int], extrapolate: int) \
+        -> Tuple[Optional[Tensor], Optional[Tensor], ]:
+    """
+    grad: (B, C, *spatial_out) tensor
+    inp: (B, C, *spatial_in) tensor
+    grid: (B, *spatial_out, D) tensor
+    bound: List{D}[int] tensor
+    interpolation: List{D}[int]
+    extrapolate: int
+    returns: (B, C, *spatial_in) tensor, (B, *spatial_out, D)
+    """
+    dim = grid.shape[-1]
+    grad_inp: Optional[Tensor] = None
+    grad_grid: Optional[Tensor] = None
+    if inp.requires_grad:
+        grad_inp = grid_push(grad, grid, inp.shape[-dim:], bound, interpolation, extrapolate)
+    if grid.requires_grad:
+        grad_grid = grid_grad(inp, grid, bound, interpolation, extrapolate)
+        # grad_grid = dot(grad_grid, grad.unsqueeze(-1), dim=1)
+        grad_grid = (grad_grid * grad.unsqueeze(-1)).sum(dim=1)
+    return grad_inp, grad_grid
+
+
+# @torch.jit.script
+def grid_push_backward(grad, inp, grid, bound: List[int],
+                       interpolation: List[int], extrapolate: int) \
+        -> Tuple[Optional[Tensor], Optional[Tensor], ]:
+    """
+    grad: (B, C, *spatial_out) tensor
+    inp: (B, C, *spatial_in) tensor
+    grid: (B, *spatial_in, D) tensor
+    bound: List{D}[int] tensor
+    interpolation: List{D}[int]
+    extrapolate: int
+    returns: (B, C, *spatial_in) tensor, (B, *spatial_in, D)
+    """
+    grad_inp: Optional[Tensor] = None
+    grad_grid: Optional[Tensor] = None
+    if inp.requires_grad:
+        grad_inp = grid_pull(grad, grid, bound, interpolation, extrapolate)
+    if grid.requires_grad:
+        grad_grid = grid_grad(grad, grid, bound, interpolation, extrapolate)
+        # grad_grid = dot(grad_grid, inp.unsqueeze(-1), dim=1)
+        grad_grid = (grad_grid * inp.unsqueeze(-1)).sum(dim=1)
+    return grad_inp, grad_grid
+
+
+# @torch.jit.script
+def grid_count_backward(grad, grid, bound: List[int],
+                       interpolation: List[int], extrapolate: int) \
+        -> Optional[Tensor]:
+    """
+    grad: (B, C, *spatial_out) tensor
+    grid: (B, *spatial_in, D) tensor
+    bound: List{D}[int] tensor
+    interpolation: List{D}[int]
+    extrapolate: int
+    returns: (B, C, *spatial_in) tensor, (B, *spatial_in, D)
+    """
+    if grid.requires_grad:
+        return grid_grad(grad, grid, bound, interpolation, extrapolate).sum(1)
+    return None
+
+
+# @torch.jit.script
+def grid_grad_backward(grad, inp, grid, bound: List[int],
+                       interpolation: List[int], extrapolate: int) \
+        -> Tuple[Optional[Tensor], Optional[Tensor]]:
+    """
+    grad: (B, C, *spatial_out, D) tensor
+    inp: (B, C, *spatial_in) tensor
+    grid: (B, *spatial_out, D) tensor
+    bound: List{D}[int] tensor
+    interpolation: List{D}[int]
+    extrapolate: int
+    returns: (B, C, *spatial_in, D) tensor, (B, *spatial_out, D)
+    """
+    dim = grid.shape[-1]
+    shape = inp.shape[-dim:]
+    grad_inp: Optional[Tensor] = None
+    grad_grid: Optional[Tensor] = None
+    if inp.requires_grad:
+        grad_inp = grid_pushgrad(grad, grid, shape, bound, interpolation, extrapolate)
+    if grid.requires_grad:
+        grad_grid = grid_hess(inp, grid, bound, interpolation, extrapolate)
+        # grad_grid = dot_multi(grad_grid, grad.unsqueeze(-1), dim=[1, -2])
+        grad_grid = (grad_grid * grad.unsqueeze(-1)).sum(dim=[1, -2])
+    return grad_inp, grad_grid

Generator/interpol/resize.py

@@ -0,0 +1,120 @@
+"""
+Resize functions (equivalent to scipy's zoom, pytorch's interpolate)
+based on grid_pull.
+"""
+__all__ = ['resize']
+
+from .api import grid_pull
+from .utils import make_list, meshgrid_ij
+from . import backend, jitfields
+import torch
+
+
+def resize(image, factor=None, shape=None, anchor='c',
+           interpolation=1, prefilter=True, **kwargs):
+    """Resize an image by a factor or to a specific shape.
+
+    Notes
+    -----
+    .. A least one of `factor` and `shape` must be specified
+    .. If `anchor in ('centers', 'edges')`, exactly one of `factor` or
+       `shape must be specified.
+    .. If `anchor in ('first', 'last')`, `factor` must be provided even
+       if `shape` is specified.
+    .. Because of rounding, it is in general not assured that
+       `resize(resize(x, f), 1/f)` returns a tensor with the same shape as x.
+
+        edges          centers          first           last
+    e - + - + - e   + - + - + - +   + - + - + - +   + - + - + - +
+    | . | . | . |   | c | . | c |   | f | . | . |   | . | . | . |
+    + _ + _ + _ +   + _ + _ + _ +   + _ + _ + _ +   + _ + _ + _ +
+    | . | . | . |   | . | . | . |   | . | . | . |   | . | . | . |
+    + _ + _ + _ +   + _ + _ + _ +   + _ + _ + _ +   + _ + _ + _ +
+    | . | . | . |   | c | . | c |   | . | . | . |   | . | . | l |
+    e _ + _ + _ e   + _ + _ + _ +   + _ + _ + _ +   + _ + _ + _ +
+
+    Parameters
+    ----------
+    image : (batch, channel, *inshape) tensor
+        Image to resize
+    factor : float or list[float], optional
+        Resizing factor
+        * > 1 : larger image <-> smaller voxels
+        * < 1 : smaller image <-> larger voxels
+    shape : (ndim,) list[int], optional
+        Output shape
+    anchor : {'centers', 'edges', 'first', 'last'} or list, default='centers'
+        * In cases 'c' and 'e', the volume shape is multiplied by the
+          zoom factor (and eventually truncated), and two anchor points
+          are used to determine the voxel size.
+        * In cases 'f' and 'l', a single anchor point is used so that
+          the voxel size is exactly divided by the zoom factor.
+          This case with an integer factor corresponds to subslicing
+          the volume (e.g., `vol[::f, ::f, ::f]`).
+        * A list of anchors (one per dimension) can also be provided.
+    interpolation : int or sequence[int], default=1
+        Interpolation order.
+    prefilter : bool, default=True
+        Apply spline pre-filter (= interpolates the input)
+
+    Returns
+    -------
+    resized : (batch, channel, *shape) tensor
+        Resized image
+
+    """
+    if backend.jitfields and jitfields.available:
+        return jitfields.resize(image, factor, shape, anchor,
+                                interpolation, prefilter, **kwargs)
+
+    factor = make_list(factor) if factor else []
+    shape = make_list(shape) if shape else []
+    anchor = make_list(anchor)
+    nb_dim = max(len(factor), len(shape), len(anchor)) or (image.dim() - 2)
+    anchor = [a[0].lower() for a in make_list(anchor, nb_dim)]
+    bck = dict(dtype=image.dtype, device=image.device)
+
+    # compute output shape
+    inshape = image.shape[-nb_dim:]
+    if factor:
+        factor = make_list(factor, nb_dim)
+    elif not shape:
+        raise ValueError('One of `factor` or `shape` must be provided')
+    if shape:
+        shape = make_list(shape, nb_dim)
+    else:
+        shape = [int(i*f) for i, f in zip(inshape, factor)]
+
+    if not factor:
+        factor = [o/i for o, i in zip(shape, inshape)]
+
+    # compute transformation grid
+    lin = []
+    for anch, f, inshp, outshp in zip(anchor, factor, inshape, shape):
+        if anch == 'c':    # centers
+            lin.append(torch.linspace(0, inshp - 1, outshp, **bck))
+        elif anch == 'e':  # edges
+            scale = inshp / outshp
+            shift = 0.5 * (scale - 1)
+            lin.append(torch.arange(0., outshp, **bck) * scale + shift)
+        elif anch == 'f':  # first voxel
+            # scale = 1/f
+            # shift = 0
+            lin.append(torch.arange(0., outshp, **bck) / f)
+        elif anch == 'l':  # last voxel
+            # scale = 1/f
+            shift = (inshp - 1) - (outshp - 1) / f
+            lin.append(torch.arange(0., outshp, **bck) / f + shift)
+        else:
+            raise ValueError('Unknown anchor {}'.format(anch))
+
+    # interpolate
+    kwargs.setdefault('bound', 'nearest')
+    kwargs.setdefault('extrapolate', True)
+    kwargs.setdefault('interpolation', interpolation)
+    kwargs.setdefault('prefilter', prefilter)
+    grid = torch.stack(meshgrid_ij(*lin), dim=-1)
+    resized = grid_pull(image, grid, **kwargs)
+
+    return resized
+

Generator/interpol/restrict.py

@@ -0,0 +1,122 @@
+__all__ = ['restrict']
+
+from .api import grid_push
+from .utils import make_list, meshgrid_ij
+from . import backend, jitfields
+import torch
+
+
+def restrict(image, factor=None, shape=None, anchor='c',
+             interpolation=1, reduce_sum=False, **kwargs):
+    """Restrict an image by a factor or to a specific shape.
+
+    Notes
+    -----
+    .. A least one of `factor` and `shape` must be specified
+    .. If `anchor in ('centers', 'edges')`, exactly one of `factor` or
+       `shape must be specified.
+    .. If `anchor in ('first', 'last')`, `factor` must be provided even
+       if `shape` is specified.
+    .. Because of rounding, it is in general not assured that
+       `resize(resize(x, f), 1/f)` returns a tensor with the same shape as x.
+
+        edges          centers          first           last
+    e - + - + - e   + - + - + - +   + - + - + - +   + - + - + - +
+    | . | . | . |   | c | . | c |   | f | . | . |   | . | . | . |
+    + _ + _ + _ +   + _ + _ + _ +   + _ + _ + _ +   + _ + _ + _ +
+    | . | . | . |   | . | . | . |   | . | . | . |   | . | . | . |
+    + _ + _ + _ +   + _ + _ + _ +   + _ + _ + _ +   + _ + _ + _ +
+    | . | . | . |   | c | . | c |   | . | . | . |   | . | . | l |
+    e _ + _ + _ e   + _ + _ + _ +   + _ + _ + _ +   + _ + _ + _ +
+
+    Parameters
+    ----------
+    image : (batch, channel, *inshape) tensor
+        Image to resize
+    factor : float or list[float], optional
+        Resizing factor
+        * > 1 : larger image <-> smaller voxels
+        * < 1 : smaller image <-> larger voxels
+    shape : (ndim,) list[int], optional
+        Output shape
+    anchor : {'centers', 'edges', 'first', 'last'} or list, default='centers'
+        * In cases 'c' and 'e', the volume shape is multiplied by the
+          zoom factor (and eventually truncated), and two anchor points
+          are used to determine the voxel size.
+        * In cases 'f' and 'l', a single anchor point is used so that
+          the voxel size is exactly divided by the zoom factor.
+          This case with an integer factor corresponds to subslicing
+          the volume (e.g., `vol[::f, ::f, ::f]`).
+        * A list of anchors (one per dimension) can also be provided.
+    interpolation : int or sequence[int], default=1
+        Interpolation order.
+    reduce_sum : bool, default=False
+        Do not normalize by the number of accumulated values per voxel
+
+    Returns
+    -------
+    restricted : (batch, channel, *shape) tensor
+        Restricted image
+
+    """
+    if backend.jitfields and jitfields.available:
+        return jitfields.restrict(image, factor, shape, anchor,
+                                  interpolation, reduce_sum, **kwargs)
+
+    factor = make_list(factor) if factor else []
+    shape = make_list(shape) if shape else []
+    anchor = make_list(anchor)
+    nb_dim = max(len(factor), len(shape), len(anchor)) or (image.dim() - 2)
+    anchor = [a[0].lower() for a in make_list(anchor, nb_dim)]
+    bck = dict(dtype=image.dtype, device=image.device)
+
+    # compute output shape
+    inshape = image.shape[-nb_dim:]
+    if factor:
+        factor = make_list(factor, nb_dim)
+    elif not shape:
+        raise ValueError('One of `factor` or `shape` must be provided')
+    if shape:
+        shape = make_list(shape, nb_dim)
+    else:
+        shape = [int(i/f) for i, f in zip(inshape, factor)]
+
+    if not factor:
+        factor = [i/o for o, i in zip(shape, inshape)]
+
+    # compute transformation grid
+    lin = []
+    fullscale = 1
+    for anch, f, inshp, outshp in zip(anchor, factor, inshape, shape):
+        if anch == 'c':    # centers
+            lin.append(torch.linspace(0, outshp - 1, inshp, **bck))
+            fullscale *= (inshp - 1) / (outshp - 1)
+        elif anch == 'e':  # edges
+            scale = outshp / inshp
+            shift = 0.5 * (scale - 1)
+            fullscale *= scale
+            lin.append(torch.arange(0., inshp, **bck) * scale + shift)
+        elif anch == 'f':  # first voxel
+            # scale = 1/f
+            # shift = 0
+            fullscale *= 1/f
+            lin.append(torch.arange(0., inshp, **bck) / f)
+        elif anch == 'l':  # last voxel
+            # scale = 1/f
+            shift = (outshp - 1) - (inshp - 1) / f
+            fullscale *= 1/f
+            lin.append(torch.arange(0., inshp, **bck) / f + shift)
+        else:
+            raise ValueError('Unknown anchor {}'.format(anch))
+
+    # scatter
+    kwargs.setdefault('bound', 'nearest')
+    kwargs.setdefault('extrapolate', True)
+    kwargs.setdefault('interpolation', interpolation)
+    kwargs.setdefault('prefilter', False)
+    grid = torch.stack(meshgrid_ij(*lin), dim=-1)
+    resized = grid_push(image, grid, shape, **kwargs)
+    if not reduce_sum:
+        resized /= fullscale
+
+    return resized

Generator/interpol/splines.py

@@ -0,0 +1,196 @@
+"""Weights and derivatives of spline orders 0 to 7."""
+import torch
+from enum import Enum
+from .jit_utils import square, cube, pow4, pow5, pow6, pow7
+
+
+class InterpolationType(Enum):
+    nearest = zeroth = 0
+    linear = first = 1
+    quadratic = second = 2
+    cubic = third = 3
+    fourth = 4
+    fifth = 5
+    sixth = 6
+    seventh = 7
+
+
+@torch.jit.script
+class Spline:
+
+    def __init__(self, order: int = 1):
+        self.order = order
+
+    def weight(self, x):
+        w = self.fastweight(x)
+        zero = torch.zeros([1], dtype=x.dtype, device=x.device)
+        w = torch.where(x.abs() >= (self.order + 1)/2, zero, w)
+        return w
+
+    def fastweight(self, x):
+        if self.order == 0:
+            return torch.ones(x.shape, dtype=x.dtype, device=x.device)
+        x = x.abs()
+        if self.order == 1:
+            return 1 - x
+        if self.order == 2:
+            x_low = 0.75 - square(x)
+            x_up = 0.5 * square(1.5 - x)
+            return torch.where(x < 0.5, x_low, x_up)
+        if self.order == 3:
+            x_low = (x * x * (x - 2.) * 3. + 4.) / 6.
+            x_up = cube(2. - x) / 6.
+            return torch.where(x < 1., x_low, x_up)
+        if self.order == 4:
+            x_low = square(x)
+            x_low = x_low * (x_low * 0.25 - 0.625) + 115. / 192.
+            x_mid = x * (x * (x * (5. - x) / 6. - 1.25) + 5./24.) + 55./96.
+            x_up = pow4(x - 2.5) / 24.
+            return torch.where(x < 0.5, x_low, torch.where(x < 1.5, x_mid, x_up))
+        if self.order == 5:
+            x_low = square(x)
+            x_low = x_low * (x_low * (0.25 - x / 12.) - 0.5) + 0.55
+            x_mid = x * (x * (x * (x * (x / 24. - 0.375) + 1.25) - 1.75) + 0.625) + 0.425
+            x_up = pow5(3 - x) / 120.
+            return torch.where(x < 1., x_low, torch.where(x < 2., x_mid, x_up))
+        if self.order == 6:
+            x_low = square(x)
+            x_low = x_low * (x_low * (7./48. - x_low/36.) - 77./192.) + 5887./11520.
+            x_mid_low = (x * (x * (x * (x * (x * (x / 48. - 7./48.) + 0.328125)
+                         - 35./288.) - 91./256.) - 7./768.) + 7861./15360.)
+            x_mid_up = (x * (x * (x * (x * (x * (7./60. - x / 120.) - 0.65625)
+                        + 133./72.) - 2.5703125) + 1267./960.) + 1379./7680.)
+            x_up = pow6(x - 3.5) / 720.
+            return torch.where(x < .5, x_low,
+                               torch.where(x < 1.5, x_mid_low,
+                                           torch.where(x < 2.5, x_mid_up, x_up)))
+        if self.order == 7:
+            x_low = square(x)
+            x_low = (x_low * (x_low * (x_low * (x / 144. - 1./36.)
+                     + 1./9.) - 1./3.) + 151./315.)
+            x_mid_low = (x * (x * (x * (x * (x * (x * (0.05 - x/240.) - 7./30.)
+                         + 0.5) - 7./18.) - 0.1) - 7./90.) + 103./210.)
+            x_mid_up = (x * (x * (x * (x * (x * (x * (x / 720. - 1./36.)
+                        + 7./30.) - 19./18.) + 49./18.) - 23./6.) + 217./90.)
+                        - 139./630.)
+            x_up = pow7(4 - x) / 5040.
+            return torch.where(x < 1., x_low,
+                               torch.where(x < 2., x_mid_low,
+                                           torch.where(x < 3., x_mid_up, x_up)))
+        raise NotImplementedError
+
+    def grad(self, x):
+        if self.order == 0:
+            return torch.zeros(x.shape, dtype=x.dtype, device=x.device)
+        g = self.fastgrad(x)
+        zero = torch.zeros([1], dtype=x.dtype, device=x.device)
+        g = torch.where(x.abs() >= (self.order + 1)/2, zero, g)
+        return g
+
+    def fastgrad(self, x):
+        if self.order == 0:
+            return torch.zeros(x.shape, dtype=x.dtype, device=x.device)
+        return self._fastgrad(x.abs()).mul(x.sign())
+
+    def _fastgrad(self, x):
+        if self.order == 1:
+            return torch.ones(x.shape, dtype=x.dtype, device=x.device)
+        if self.order == 2:
+            return torch.where(x < 0.5, -2*x, x - 1.5)
+        if self.order == 3:
+            g_low = x * (x * 1.5 - 2)
+            g_up = -0.5 * square(2 - x)
+            return torch.where(x < 1, g_low, g_up)
+        if self.order == 4:
+            g_low = x * (square(x) - 1.25)
+            g_mid = x * (x * (x * (-2./3.) + 2.5) - 2.5) + 5./24.
+            g_up = cube(2. * x - 5.) / 48.
+            return torch.where(x < 0.5, g_low,
+                               torch.where(x < 1.5, g_mid, g_up))
+        if self.order == 5:
+            g_low = x * (x * (x * (x * (-5./12.) + 1.)) - 1.)
+            g_mid = x * (x * (x * (x * (5./24.) - 1.5) + 3.75) - 3.5) + 0.625
+            g_up = pow4(x - 3.) / (-24.)
+            return torch.where(x < 1, g_low,
+                               torch.where(x < 2, g_mid, g_up))
+        if self.order == 6:
+            g_low = square(x)
+            g_low = x * (g_low * (7./12.) - square(g_low) / 6. - 77./96.)
+            g_mid_low = (x * (x * (x * (x * (x * 0.125 - 35./48.) + 1.3125)
+                         - 35./96.) - 0.7109375) - 7./768.)
+            g_mid_up = (x * (x * (x * (x * (x / (-20.) + 7./12.) - 2.625)
+                        + 133./24.) - 5.140625) + 1267./960.)
+            g_up = pow5(2*x - 7) / 3840.
+            return torch.where(x < 0.5, g_low,
+                               torch.where(x < 1.5, g_mid_low,
+                                           torch.where(x < 2.5, g_mid_up,
+                                                       g_up)))
+        if self.order == 7:
+            g_low = square(x)
+            g_low = x * (g_low * (g_low * (x * (7./144.) - 1./6.) + 4./9.) - 2./3.)
+            g_mid_low = (x * (x * (x * (x * (x * (x * (-7./240.) + 3./10.)
+                         - 7./6.) + 2.) - 7./6.) - 1./5.) - 7./90.)
+            g_mid_up = (x * (x * (x * (x * (x * (x * (7./720.) - 1./6.)
+                        + 7./6.) - 38./9.) + 49./6.) - 23./3.) + 217./90.)
+            g_up = pow6(x - 4) / (-720.)
+            return torch.where(x < 1, g_low,
+                               torch.where(x < 2, g_mid_low,
+                                           torch.where(x < 3, g_mid_up, g_up)))
+        raise NotImplementedError
+
+    def hess(self, x):
+        if self.order == 0:
+            return torch.zeros(x.shape, dtype=x.dtype, device=x.device)
+        h = self.fasthess(x)
+        zero = torch.zeros([1], dtype=x.dtype, device=x.device)
+        h = torch.where(x.abs() >= (self.order + 1)/2, zero, h)
+        return h
+
+    def fasthess(self, x):
+        if self.order in (0, 1):
+            return torch.zeros(x.shape, dtype=x.dtype, device=x.device)
+        x = x.abs()
+        if self.order == 2:
+            one = torch.ones([1], dtype=x.dtype, device=x.device)
+            return torch.where(x < 0.5, -2 * one, one)
+        if self.order == 3:
+            return torch.where(x < 1, 3. * x - 2., 2. - x)
+        if self.order == 4:
+            return torch.where(x < 0.5, 3. * square(x) - 1.25,
+                               torch.where(x < 1.5, x * (-2. * x + 5.) - 2.5,
+                                           square(2. * x - 5.) / 8.))
+        if self.order == 5:
+            h_low = square(x)
+            h_low = - h_low * (x * (5./3.) - 3.) - 1.
+            h_mid = x * (x * (x * (5./6.) - 9./2.) + 15./2.) - 7./2.
+            h_up = 9./2. - x * (x * (x/6. - 3./2.) + 9./2.)
+            return torch.where(x < 1, h_low,
+                               torch.where(x < 2, h_mid, h_up))
+        if self.order == 6:
+            h_low = square(x)
+            h_low = - h_low * (h_low * (5./6) - 7./4.) - 77./96.
+            h_mid_low = (x * (x * (x * (x * (5./8.) - 35./12.) + 63./16.)
+                         - 35./48.) - 91./128.)
+            h_mid_up = -(x * (x * (x * (x/4. - 7./3.) + 63./8.) - 133./12.)
+                         + 329./64.)
+            h_up = (x * (x * (x * (x/24. - 7./12.) + 49./16.) - 343./48.)
+                    + 2401./384.)
+            return torch.where(x < 0.5, h_low,
+                               torch.where(x < 1.5, h_mid_low,
+                                           torch.where(x < 2.5, h_mid_up,
+                                                       h_up)))
+        if self.order == 7:
+            h_low = square(x)
+            h_low = h_low * (h_low*(x * (7./24.) - 5./6.) + 4./3.) - 2./3.
+            h_mid_low = - (x * (x * (x * (x * (x * (7./40.) - 3./2.) + 14./3.)
+                           - 6.) + 7./3.) + 1./5.)
+            h_mid_up = (x * (x * (x * (x * (x * (7./120.) - 5./6.) + 14./3.)
+                        - 38./3.) + 49./3.) - 23./3.)
+            h_up = - (x * (x * (x * (x * (x/120. - 1./6.) + 4./3.) - 16./3.)
+                      + 32./3.) - 128./15.)
+            return torch.where(x < 1, h_low,
+                               torch.where(x < 2, h_mid_low,
+                                           torch.where(x < 3, h_mid_up,
+                                                       h_up)))
+        raise NotImplementedError
+

Generator/interpol/tests/test_gradcheck_pushpull.py

@@ -0,0 +1,125 @@
+import torch
+from torch.autograd import gradcheck
+from interpol import grid_pull, grid_push, grid_count, grid_grad, add_identity_grid_
+import pytest
+import inspect
+
+# global parameters
+dtype = torch.double        # data type (double advised to check gradients)
+shape1 = 3                  # size along each dimension
+extrapolate = True
+
+if hasattr(torch, 'use_deterministic_algorithms'):
+    torch.use_deterministic_algorithms(True)
+kwargs = dict(rtol=1., raise_exception=True)
+if 'check_undefined_grad' in inspect.signature(gradcheck).parameters:
+    kwargs['check_undefined_grad'] = False
+if 'nondet_tol' in inspect.signature(gradcheck).parameters:
+    kwargs['nondet_tol'] = 1e-3
+
+# parameters
+devices = [('cpu', 1)]
+if torch.backends.openmp.is_available() or torch.backends.mkl.is_available():
+    print('parallel backend available')
+    devices.append(('cpu', 10))
+if torch.cuda.is_available():
+    print('cuda backend available')
+    devices.append('cuda')
+
+dims = [1, 2, 3]
+bounds = list(range(7))
+order_bounds = []
+for o in range(3):
+    for b in bounds:
+        order_bounds += [(o, b)]
+for o in range(3, 8):
+    order_bounds += [(o, 3)]  # only test dc2 for order > 2
+
+
+def make_data(shape, device, dtype):
+    grid = torch.randn([2, *shape, len(shape)], device=device, dtype=dtype)
+    grid = add_identity_grid_(grid)
+    vol = torch.randn((2, 1,) + shape, device=device, dtype=dtype)
+    return vol, grid
+
+
+def init_device(device):
+    if isinstance(device, (list, tuple)):
+        device, param = device
+    else:
+        param = 1 if device == 'cpu' else 0
+    if device == 'cuda':
+        torch.cuda.set_device(param)
+        torch.cuda.init()
+        try:
+            torch.cuda.empty_cache()
+        except RuntimeError:
+            pass
+        device = '{}:{}'.format(device, param)
+    else:
+        assert device == 'cpu'
+        torch.set_num_threads(param)
+    return torch.device(device)
+
+
+@pytest.mark.parametrize("device", devices)
+@pytest.mark.parametrize("dim", dims)
+# @pytest.mark.parametrize("bound", bounds)
+# @pytest.mark.parametrize("interpolation", orders)
+@pytest.mark.parametrize("interpolation,bound", order_bounds)
+def test_gradcheck_grad(device, dim, bound, interpolation):
+    print(f'grad_{dim}d({interpolation}, {bound}) on {device}')
+    device = init_device(device)
+    shape = (shape1,) * dim
+    vol, grid = make_data(shape, device, dtype)
+    vol.requires_grad = True
+    grid.requires_grad = True
+    assert gradcheck(grid_grad, (vol, grid, interpolation, bound, extrapolate),
+                     **kwargs)
+
+
+@pytest.mark.parametrize("device", devices)
+@pytest.mark.parametrize("dim", dims)
+# @pytest.mark.parametrize("bound", bounds)
+# @pytest.mark.parametrize("interpolation", orders)
+@pytest.mark.parametrize("interpolation,bound", order_bounds)
+def test_gradcheck_pull(device, dim, bound, interpolation):
+    print(f'pull_{dim}d({interpolation}, {bound}) on {device}')
+    device = init_device(device)
+    shape = (shape1,) * dim
+    vol, grid = make_data(shape, device, dtype)
+    vol.requires_grad = True
+    grid.requires_grad = True
+    assert gradcheck(grid_pull, (vol, grid, interpolation, bound, extrapolate),
+                     **kwargs)
+
+
+@pytest.mark.parametrize("device", devices)
+@pytest.mark.parametrize("dim", dims)
+# @pytest.mark.parametrize("bound", bounds)
+# @pytest.mark.parametrize("interpolation", orders)
+@pytest.mark.parametrize("interpolation,bound", order_bounds)
+def test_gradcheck_push(device, dim, bound, interpolation):
+    print(f'push_{dim}d({interpolation}, {bound}) on {device}')
+    device = init_device(device)
+    shape = (shape1,) * dim
+    vol, grid = make_data(shape, device, dtype)
+    vol.requires_grad = True
+    grid.requires_grad = True
+    assert gradcheck(grid_push, (vol, grid, shape, interpolation, bound, extrapolate),
+                     **kwargs)
+
+
+@pytest.mark.parametrize("device", devices)
+@pytest.mark.parametrize("dim", dims)
+# @pytest.mark.parametrize("bound", bounds)
+# @pytest.mark.parametrize("interpolation", orders)
+@pytest.mark.parametrize("interpolation,bound", order_bounds)
+def test_gradcheck_count(device, dim, bound, interpolation):
+    print(f'count_{dim}d({interpolation}, {bound}) on {device}')
+    device = init_device(device)
+    shape = (shape1,) * dim
+    _, grid = make_data(shape, device, dtype)
+    grid.requires_grad = True
+    assert gradcheck(grid_count, (grid, shape, interpolation, bound, extrapolate),
+                     **kwargs)

Generator/interpol/utils.py

@@ -0,0 +1,176 @@
+import torch
+
+
+def fake_decorator(*a, **k):
+    if len(a) == 1 and not k:
+        return a[0]
+    else:
+        return fake_decorator
+
+
+def make_list(x, n=None, **kwargs):
+    """Ensure that the input  is a list (of a given size)
+
+    Parameters
+    ----------
+    x : list or tuple or scalar
+        Input object
+    n : int, optional
+        Required length
+    default : scalar, optional
+        Value to right-pad with. Use last value of the input by default.
+
+    Returns
+    -------
+    x : list
+    """
+    if not isinstance(x, (list, tuple)):
+        x = [x]
+    x = list(x)
+    if n and len(x) < n:
+        default = kwargs.get('default', x[-1])
+        x = x + [default] * max(0, n - len(x))
+    return x
+
+
+def expanded_shape(*shapes, side='left'):
+    """Expand input shapes according to broadcasting rules
+
+    Parameters
+    ----------
+    *shapes : sequence[int]
+        Input shapes
+    side : {'left', 'right'}, default='left'
+        Side to add singleton dimensions.
+
+    Returns
+    -------
+    shape : tuple[int]
+        Output shape
+
+    Raises
+    ------
+    ValueError
+        If shapes are not compatible for broadcast.
+
+    """
+    def error(s0, s1):
+        raise ValueError('Incompatible shapes for broadcasting: {} and {}.'
+                         .format(s0, s1))
+
+    # 1. nb dimensions
+    nb_dim = 0
+    for shape in shapes:
+        nb_dim = max(nb_dim, len(shape))
+
+    # 2. enumerate
+    shape = [1] * nb_dim
+    for i, shape1 in enumerate(shapes):
+        pad_size = nb_dim - len(shape1)
+        ones = [1] * pad_size
+        if side == 'left':
+            shape1 = [*ones, *shape1]
+        else:
+            shape1 = [*shape1, *ones]
+        shape = [max(s0, s1) if s0 == 1 or s1 == 1 or s0 == s1
+                 else error(s0, s1) for s0, s1 in zip(shape, shape1)]
+
+    return tuple(shape)
+
+
+def matvec(mat, vec, out=None):
+    """Matrix-vector product (supports broadcasting)
+
+    Parameters
+    ----------
+    mat : (..., M, N) tensor
+        Input matrix.
+    vec : (..., N) tensor
+        Input vector.
+    out : (..., M) tensor, optional
+        Placeholder for the output tensor.
+
+    Returns
+    -------
+    mv : (..., M) tensor
+        Matrix vector product of the inputs
+
+    """
+    vec = vec[..., None]
+    if out is not None:
+        out = out[..., None]
+
+    mv = torch.matmul(mat, vec, out=out)
+    mv = mv[..., 0]
+    if out is not None:
+        out = out[..., 0]
+
+    return mv
+
+
+def _compare_versions(version1, mode, version2):
+    for v1, v2 in zip(version1, version2):
+        if mode in ('gt', '>'):
+            if v1 > v2:
+                return True
+            elif v1 < v2:
+                return False
+        elif mode in ('ge', '>='):
+            if v1 > v2:
+                return True
+            elif v1 < v2:
+                return False
+        elif mode in ('lt', '<'):
+            if v1 < v2:
+                return True
+            elif v1 > v2:
+                return False
+        elif mode in ('le', '<='):
+            if v1 < v2:
+                return True
+            elif v1 > v2:
+                return False
+    if mode in ('gt', 'lt', '>', '<'):
+        return False
+    else:
+        return True
+
+
+def torch_version(mode, version):
+    """Check torch version
+
+    Parameters
+    ----------
+    mode : {'<', '<=', '>', '>='}
+    version : tuple[int]
+
+    Returns
+    -------
+    True if "torch.version <mode> version"
+
+    """
+    current_version, *cuda_variant = torch.__version__.split('+')
+    major, minor, patch, *_ = current_version.split('.')
+    # strip alpha tags
+    for x in 'abcdefghijklmnopqrstuvwxy':
+        if x in patch:
+            patch = patch[:patch.index(x)]
+    current_version = (int(major), int(minor), int(patch))
+    version = make_list(version)
+    return _compare_versions(current_version, mode, version)
+
+
+if torch_version('>=', (1, 10)):
+    meshgrid_ij = lambda *x: torch.meshgrid(*x, indexing='ij')
+    meshgrid_xy = lambda *x: torch.meshgrid(*x, indexing='xy')
+else:
+    meshgrid_ij = lambda *x: torch.meshgrid(*x)
+    def meshgrid_xy(*x):
+        grid = list(torch.meshgrid(*x))
+        if len(grid) > 1:
+            grid[0] = grid[0].transpose(0, 1)
+            grid[1] = grid[1].transpose(0, 1)
+        return grid
+
+
+meshgrid = meshgrid_ij

Generator/utils.py

@@ -0,0 +1,669 @@
+import os
+import numpy as np
+import nibabel as nib
+
+import torch
+from torch.nn.functional import conv3d
+from torch.utils.data import Dataset
+
+from scipy.io.matlab import loadmat
+
+
+import time, datetime
+
+from ShapeID.DiffEqs.adjoint import odeint_adjoint as odeint
+from ShapeID.perlin3d import generate_velocity_3d , generate_shape_3d
+
+
+class ConcatDataset(Dataset):
+    def __init__(self,dataset_list, probs=None):
+        self.datasets = dataset_list
+        self.probs = probs if probs else [1/len(self.datasets)] * len(self.datasets)
+
+    def __getitem__(self, i):
+        chosen_dataset = np.random.choice(self.datasets, 1, p=self.probs)[0]
+        i = i % len(chosen_dataset)
+        return chosen_dataset[i]
+
+    def __len__(self):
+        return  max(len(d) for d in self.datasets)
+
+
+
+# Prepare generator
+def resolution_sampler(low_res_only = False):
+    
+    if low_res_only:
+        r = (np.random.rand() * 0.5) + 0.5 # in [0.5, 1]
+    else:
+        r = np.random.rand() # in [0, 1]
+
+    if r < 0.25: # 1mm isotropic
+        resolution = np.array([1.0, 1.0, 1.0])
+        thickness = np.array([1.0, 1.0, 1.0])
+    elif r < 0.5: # clinical (low-res in one dimension)
+        resolution = np.array([1.0, 1.0, 1.0])
+        thickness = np.array([1.0, 1.0, 1.0])
+        idx = np.random.randint(3)
+        resolution[idx] = 2.5 + 6 * np.random.rand()
+        thickness[idx] = np.min([resolution[idx], 4.0 + 2.0 * np.random.rand()])
+    elif r < 0.75:  # low-field: stock sequences (always axial)
+        resolution = np.array([1.3, 1.3, 4.8]) + 0.4 * np.random.rand(3)
+        thickness = resolution.copy()
+    else: # low-field: isotropic-ish (also good for scouts)
+        resolution = 2.0 + 3.0 * np.random.rand(3)
+        thickness = resolution.copy()
+
+    return resolution, thickness
+
+
+#####################################
+############ Utility Func ###########
+#####################################
+
+
+def binarize(p, thres):
+    # TODO: what is the optimal thresholding strategy?
+    thres = thres * p.max()
+
+    bin = p.clone()
+    bin[p < thres] = 0.
+    bin[p >= thres] = 1. 
+    return bin
+
+def make_gaussian_kernel(sigma, device):
+
+    sl = int(np.ceil(3 * sigma))
+    ts = torch.linspace(-sl, sl, 2*sl+1, dtype=torch.float, device=device)
+    gauss = torch.exp((-(ts / sigma)**2 / 2))
+    kernel = gauss / gauss.sum()
+
+    return kernel
+
+def gaussian_blur_3d(input, stds, device):
+    blurred = input[None, None, :, :, :]
+    if stds[0]>0:
+        kx = make_gaussian_kernel(stds[0], device=device)
+        blurred = conv3d(blurred, kx[None, None, :, None, None], stride=1, padding=(len(kx) // 2, 0, 0))
+    if stds[1]>0:
+        ky = make_gaussian_kernel(stds[1], device=device)
+        blurred = conv3d(blurred, ky[None, None, None, :, None], stride=1, padding=(0, len(ky) // 2, 0))
+    if stds[2]>0:
+        kz = make_gaussian_kernel(stds[2], device=device)
+        blurred = conv3d(blurred, kz[None, None, None, None, :], stride=1, padding=(0, 0, len(kz) // 2))
+    return torch.squeeze(blurred)
+
+
+
+#####################################
+######### Deformation Func ##########
+#####################################
+
+def make_affine_matrix(rot, sh, s):
+    Rx = np.array([[1, 0, 0], [0, np.cos(rot[0]), -np.sin(rot[0])], [0, np.sin(rot[0]), np.cos(rot[0])]])
+    Ry = np.array([[np.cos(rot[1]), 0, np.sin(rot[1])], [0, 1, 0], [-np.sin(rot[1]), 0, np.cos(rot[1])]])
+    Rz = np.array([[np.cos(rot[2]), -np.sin(rot[2]), 0], [np.sin(rot[2]), np.cos(rot[2]), 0], [0, 0, 1]])
+
+    SHx = np.array([[1, 0, 0], [sh[1], 1, 0], [sh[2], 0, 1]])
+    SHy = np.array([[1, sh[0], 0], [0, 1, 0], [0, sh[2], 1]])
+    SHz = np.array([[1, 0, sh[0]], [0, 1, sh[1]], [0, 0, 1]])
+
+    A = SHx @ SHy @ SHz @ Rx @ Ry @ Rz
+    A[0, :] = A[0, :] * s[0]
+    A[1, :] = A[1, :] * s[1]
+    A[2, :] = A[2, :] * s[2]
+
+    return A
+
+
+def fast_3D_interp_torch(X, II, JJ, KK, mode='linear', default_value_linear=0.0):
+
+    if II is None: 
+        return X
+ 
+    if mode=='nearest':
+        IIr = torch.round(II).long()
+        JJr = torch.round(JJ).long()
+        KKr = torch.round(KK).long()
+        IIr[IIr < 0] = 0
+        JJr[JJr < 0] = 0
+        KKr[KKr < 0] = 0
+        IIr[IIr > (X.shape[0] - 1)] = (X.shape[0] - 1)
+        JJr[JJr > (X.shape[1] - 1)] = (X.shape[1] - 1)
+        KKr[KKr > (X.shape[2] - 1)] = (X.shape[2] - 1)
+        if len(X.shape)==3:
+            X = X[..., None] 
+        Y = X[IIr, JJr, KKr]
+        if Y.shape[3] == 1:
+            Y = Y[:, :, :, 0]
+
+    elif mode=='linear':
+        ok = (II>0) & (JJ>0) & (KK>0) & (II<=X.shape[0]-1) & (JJ<=X.shape[1]-1) & (KK<=X.shape[2]-1)
+        
+        IIv = II[ok]
+        JJv = JJ[ok]
+        KKv = KK[ok]
+
+        fx = torch.floor(IIv).long()
+        cx = fx + 1
+        cx[cx > (X.shape[0] - 1)] = (X.shape[0] - 1)
+        wcx = (IIv - fx)[..., None]
+        wfx = 1 - wcx
+
+        fy = torch.floor(JJv).long()
+        cy = fy + 1
+        cy[cy > (X.shape[1] - 1)] = (X.shape[1] - 1)
+        wcy = (JJv - fy)[..., None]
+        wfy = 1 - wcy
+
+        fz = torch.floor(KKv).long()
+        cz = fz + 1
+        cz[cz > (X.shape[2] - 1)] = (X.shape[2] - 1)
+        wcz = (KKv - fz)[..., None]
+        wfz = 1 - wcz
+
+        if len(X.shape)==3:
+            X = X[..., None] 
+        
+        c000 = X[fx, fy, fz]
+        c100 = X[cx, fy, fz]
+        c010 = X[fx, cy, fz]
+        c110 = X[cx, cy, fz]
+        c001 = X[fx, fy, cz]
+        c101 = X[cx, fy, cz]
+        c011 = X[fx, cy, cz]
+        c111 = X[cx, cy, cz]
+
+        c00 = c000 * wfx + c100 * wcx
+        c01 = c001 * wfx + c101 * wcx
+        c10 = c010 * wfx + c110 * wcx
+        c11 = c011 * wfx + c111 * wcx
+
+        c0 = c00 * wfy + c10 * wcy
+        c1 = c01 * wfy + c11 * wcy
+
+        c = c0 * wfz + c1 * wcz
+
+        Y = torch.zeros([*II.shape, X.shape[3]], device=X.device) 
+        Y[ok] = c.float()
+        Y[~ok] = default_value_linear   
+
+        if Y.shape[-1]==1:
+            Y = Y[...,0] 
+    else:
+        raise Exception('mode must be linear or nearest')
+    
+    return Y
+
+
+
+def myzoom_torch(X, factor, aff=None):
+
+    if len(X.shape)==3:
+        X = X[..., None]
+
+    delta = (1.0 - factor) / (2.0 * factor)
+    newsize = np.round(X.shape[:-1] * factor).astype(int)
+
+    vx = torch.arange(delta[0], delta[0] + newsize[0] / factor[0], 1 / factor[0], dtype=torch.float, device=X.device)[:newsize[0]]
+    vy = torch.arange(delta[1], delta[1] + newsize[1] / factor[1], 1 / factor[1], dtype=torch.float, device=X.device)[:newsize[1]]
+    vz = torch.arange(delta[2], delta[2] + newsize[2] / factor[2], 1 / factor[2], dtype=torch.float, device=X.device)[:newsize[2]]
+
+    vx[vx < 0] = 0
+    vy[vy < 0] = 0
+    vz[vz < 0] = 0
+    vx[vx > (X.shape[0]-1)] = (X.shape[0]-1)
+    vy[vy > (X.shape[1] - 1)] = (X.shape[1] - 1)
+    vz[vz > (X.shape[2] - 1)] = (X.shape[2] - 1)
+
+    fx = torch.floor(vx).int()
+    cx = fx + 1
+    cx[cx > (X.shape[0]-1)] = (X.shape[0]-1)
+    wcx = (vx - fx) 
+    wfx = 1 - wcx
+
+    fy = torch.floor(vy).int()
+    cy = fy + 1
+    cy[cy > (X.shape[1]-1)] = (X.shape[1]-1)
+    wcy = (vy - fy) 
+    wfy = 1 - wcy
+
+    fz = torch.floor(vz).int()
+    cz = fz + 1
+    cz[cz > (X.shape[2]-1)] = (X.shape[2]-1)
+    wcz = (vz - fz) 
+    wfz = 1 - wcz
+
+    Y = torch.zeros([newsize[0], newsize[1], newsize[2], X.shape[3]], dtype=torch.float, device=X.device) 
+
+    tmp1 = torch.zeros([newsize[0], X.shape[1], X.shape[2], X.shape[3]], dtype=torch.float, device=X.device)
+    for i in range(newsize[0]):
+        tmp1[i, :, :] = wfx[i] * X[fx[i], :, :] +  wcx[i] * X[cx[i], :, :]
+    tmp2 = torch.zeros([newsize[0], newsize[1], X.shape[2], X.shape[3]], dtype=torch.float, device=X.device)
+    for j in range(newsize[1]):
+        tmp2[:, j, :] = wfy[j] * tmp1[:, fy[j], :] +  wcy[j] * tmp1[:, cy[j], :]
+    for k in range(newsize[2]):
+        Y[:, :, k] = wfz[k] * tmp2[:, :, fz[k]] +  wcz[k] * tmp2[:, :, cz[k]]
+
+    if Y.shape[3] == 1:
+        Y = Y[:,:,:, 0]
+
+    if aff is not None:
+        aff_new = aff.copy() 
+        aff_new[:-1] = aff_new[:-1] / factor
+        aff_new[:-1, -1] = aff_new[:-1, -1] - aff[:-1, :-1] @ (0.5 - 0.5 / (factor * np.ones(3)))
+        return Y, aff_new
+    else:
+        return Y
+    
+
+
+
+#####################################
+############ Reading Func ###########
+#####################################
+
+def read_image(file_name):
+    img = nib.load(file_name) 
+    aff = img.affine
+    res = np.sqrt(np.sum(abs(aff[:-1, :-1]), axis=0)) 
+    return img, aff, res
+
+def deform_image(I, deform_dict, device, default_value_linear_mode=None, deform_mode = 'linear'):
+    if I is None:
+        return I
+    
+    [xx2, yy2, zz2, x1, y1, z1, x2, y2, z2] = deform_dict['grid']
+
+    if not isinstance(I, torch.Tensor):
+        I = torch.squeeze(torch.tensor(I.get_fdata()[x1:x2, y1:y2, z1:z2].astype(float), dtype=torch.float, device=device))
+    else:
+        I = torch.squeeze(I[x1:x2, y1:y2, z1:z2].astype(float), dtype=torch.float, device=device)
+    I = torch.nan_to_num(I) 
+    
+    if default_value_linear_mode is not None:
+        if default_value_linear_mode == 'max':
+            default_value_linear = torch.max(I)
+        else:
+            raise ValueError('Not support default_value_linear_mode:', default_value_linear_mode)
+    else:
+        default_value_linear = 0.
+    Idef = fast_3D_interp_torch(I, xx2, yy2, zz2, deform_mode, default_value_linear) 
+    
+    return Idef
+
+
+def read_and_deform(file_name, dtype, deform_dict, device, mask, default_value_linear_mode=None, deform_mode = 'linear', mean = 0., scale = 1.):
+    [xx2, yy2, zz2, x1, y1, z1, x2, y2, z2] = deform_dict['grid']
+
+    try:
+        Iimg = nib.load(file_name)  
+    except:
+        Iimg = nib.load(file_name + '.gz')
+    res = np.sqrt(np.sum(abs(Iimg.affine[:-1, :-1]), axis=0))
+    I = torch.squeeze(torch.tensor(Iimg.get_fdata()[x1:x2, y1:y2, z1:z2].astype(float), dtype=dtype, device=device))
+    I = torch.nan_to_num(I) 
+
+    I -= mean
+    I /= scale
+
+    if mask is not None:
+        I[mask == 0] = 0
+
+    if default_value_linear_mode is not None:
+        if default_value_linear_mode == 'max':
+            default_value_linear = torch.max(I)
+        else:
+            raise ValueError('Not support default_value_linear_mode:', default_value_linear_mode)
+    else:
+        default_value_linear = 0.
+    Idef = fast_3D_interp_torch(I, xx2, yy2, zz2, deform_mode, default_value_linear) 
+    return Idef, res
+
+
+def read_and_deform_image(exist_keys, task_name, file_name, setups, deform_dict, device, mask, **kwargs):
+    Idef, _ = read_and_deform(file_name, torch.float, deform_dict, device, mask) 
+    Idef -= torch.min(Idef)
+    Idef /= torch.max(Idef)
+    if setups['flip']: 
+        Idef = torch.flip(Idef, [0]) 
+    update_dict = {task_name: Idef[None]}
+    
+    if os.path.isfile(file_name[:-4] + '.defacingmask.nii'):
+        Idef_DM, _ = read_and_deform(file_name[:-4] + '.defacingmask.nii', torch.float, deform_dict, device, mask) 
+        Idef_DM = torch.clamp(Idef_DM, min = 0.) 
+        Idef_DM /= torch.max(Idef_DM) 
+        if setups['flip']: 
+            Idef = torch.flip(Idef_DM, [0]) 
+        update_dict.update({task_name + '_DM': Idef_DM[None]}) 
+    #if not 'brain_mask' in exist_keys:
+    #    mask = torch.ones_like(Idef)
+    #    mask[Idef <= 0.] = 0.  
+    #    update_dict.update({'brain_mask': mask[None]}) 
+    return update_dict
+
+def read_and_deform_CT(exist_keys, task_name, file_name, setups, deform_dict, device, mask, **kwargs):
+    Idef, _ = read_and_deform(file_name, torch.float, deform_dict, device, mask, scale = 1000)
+    #Idef = torch.clamp(Idef, min = 0., max = 80.) # No clamping for inference/GT
+    #Idef /= torch.max(Idef)
+    if setups['flip']: 
+        Idef = torch.flip(Idef, [0]) 
+    update_dict = {'CT': Idef[None]}
+    
+    if os.path.isfile(file_name[:-4] + '.defacingmask.nii'):
+        Idef_DM, _ = read_and_deform(file_name[:-4] + '.defacingmask.nii', torch.float, deform_dict, device, mask) 
+        Idef_DM = torch.clamp(Idef_DM, min = 0.) 
+        Idef_DM /= torch.max(Idef_DM) 
+        if setups['flip']: 
+            Idef = torch.flip(Idef_DM, [0]) 
+        update_dict.update({task_name + '_DM': Idef_DM[None]}) 
+    #if not 'brain_mask' in exist_keys:
+    #    mask = torch.ones_like(Idef)
+    #    mask[Idef <= 0.] = 0. 
+    #    update_dict.update({'brain_mask': mask[None]}) 
+    return update_dict
+
+def read_and_deform_distance(exist_keys, task_name, file_names, setups, deform_dict, device, mask, cfg, **kwargs):
+    [lp_dist_map, lw_dist_map, rp_dist_map, rw_dist_map] = file_names
+    
+    
+    lp, _ = read_and_deform(lp_dist_map, torch.float, deform_dict, device, mask, default_value_linear_mode = 'max', mean = 128., scale = 20) 
+    lw, _ = read_and_deform(lw_dist_map, torch.float, deform_dict, device, mask, default_value_linear_mode = 'max', mean = 128., scale = 20) 
+
+    if mask is not None: # left_hemis_only
+        Idef = torch.stack([lp, lw], dim = 0)  
+    else: 
+        rp, _ = read_and_deform(rp_dist_map, torch.float, deform_dict, device, mask, default_value_linear_mode = 'max', mean = 128., scale = 20) 
+        rw, _ = read_and_deform(rw_dist_map, torch.float, deform_dict, device, mask, default_value_linear_mode = 'max', mean = 128., scale = 20) 
+
+        if setups['flip']: 
+            aux = torch.flip(lp, [0])
+            lp = torch.flip(rp, [0])
+            rp = aux
+            aux = torch.flip(lw, [0])
+            lw = torch.flip(rw, [0])
+            rw = aux
+
+        Idef = torch.stack([lp, lw, rp, rw], dim = 0)
+
+    Idef /= deform_dict['scaling_factor_distances']
+    Idef = torch.clamp(Idef, min=-cfg.max_surf_distance, max=cfg.max_surf_distance)
+    
+    return {'distance':  Idef} 
+    
+def read_and_deform_segmentation(exist_keys, task_name, file_name, setups, deform_dict, device, mask, cfg, onehotmatrix, lut, vflip, **kwargs):
+    [xx2, yy2, zz2, x1, y1, z1, x2, y2, z2] = deform_dict['grid']
+
+    Simg = nib.load(file_name)
+    S = torch.squeeze(torch.tensor(Simg.get_fdata()[x1:x2, y1:y2, z1:z2].astype(int), dtype=torch.int, device=device))
+
+    if mask is not None:
+        S[mask == 0] = 0
+
+    Sdef = fast_3D_interp_torch(S,  xx2, yy2, zz2, 'nearest')
+    if cfg.generator.deform_one_hots:
+        Sonehot = onehotmatrix[lut[S.long()]]
+        Sdef_OneHot = fast_3D_interp_torch(Sonehot, xx2, yy2, zz2)
+    else:
+        Sdef_OneHot = onehotmatrix[lut[Sdef.long()]]
+        
+    if setups['flip']: 
+        #Sdef = torch.flip(Sdef, [0])   
+        Sdef_OneHot = torch.flip(Sdef_OneHot, [0])[:, :, :, vflip]
+
+    # prepare for input
+    Sdef_OneHot = Sdef_OneHot.permute([3, 0, 1, 2])
+
+    #update_dict = {'label': Sdef[None], 'segmentation': Sdef_OneHot}
+    update_dict = {'segmentation': Sdef_OneHot}
+
+    #if not 'brain_mask' in exist_keys:
+    #    mask = torch.ones_like(Sdef)
+    #    mask[Sdef <= 0.] = 0. 
+    #    update_dict.update({'brain_mask': mask[None]}) 
+    return update_dict
+
+
+
+def read_and_deform_pathology(exist_keys, task_name, file_name, setups, deform_dict, device, mask = None, 
+                              augment = False, pde_func = None, t = None, 
+                              shape_gen_args = None, thres = 0., **kwargs):
+    # NOTE does not support left_hemis for now
+
+    [xx2, yy2, zz2, x1, y1, z1, x2, y2, z2] = deform_dict['grid']
+
+    if file_name is None:
+        return {'pathology': torch.zeros(xx2.shape)[None].to(device), 'pathology_prob': torch.zeros(xx2.shape)[None].to(device)}
+    
+    if file_name == 'random_shape': # generate random shape
+        percentile = np.random.uniform(shape_gen_args.mask_percentile_min, shape_gen_args.mask_percentile_max)
+        _, Pdef = generate_shape_3d(xx2.shape, shape_gen_args.perlin_res, percentile, device)   
+    else: # read from existing shape
+        Pdef, _ = read_and_deform(file_name, torch.float, deform_dict, device)  
+
+    if augment:
+        Pdef = augment_pathology(Pdef, pde_func, t, shape_gen_args, device) 
+
+    #if setups['flip']: # flipping should happen after P has been encoded
+    #    Pdef = torch.flip(Pdef, [0]) 
+
+    P = binarize(Pdef, thres)
+    if P.mean() <= shape_gen_args.pathol_tol:
+        return {'pathology': torch.zeros(xx2.shape)[None].to(device), 'pathology_prob': torch.zeros(xx2.shape)[None].to(device)}
+    #print('process', P.mean(), shape_gen_args.pathol_tol)
+
+    return {'pathology': P[None], 'pathology_prob': Pdef[None]}
+
+
+def read_and_deform_registration(exist_keys, task_name, file_names, setups, deform_dict, device, mask, **kwargs):
+    [mni_reg_x, mni_reg_y, mni_reg_z] = file_names
+    regx, _ = read_and_deform(mni_reg_x, torch.float, deform_dict, device, mask, scale = 10000) 
+    regy, _ = read_and_deform(mni_reg_y, torch.float, deform_dict, device, mask, scale = 10000) 
+    regz, _ = read_and_deform(mni_reg_z, torch.float, deform_dict, device, mask, scale = 10000)  
+
+    if setups['flip']: 
+        regx = -torch.flip(regx, [0]) # NOTE: careful with switching sign
+        regy = torch.flip(regy, [0])
+        regz = torch.flip(regz, [0])
+
+    Idef = torch.stack([regx, regy, regz], dim = 0) 
+    
+    return {'registration':  Idef}
+
+def read_and_deform_bias_field(exist_keys, task_name, file_name, setups, deform_dict, device, mask, **kwargs):
+    Idef, _ = read_and_deform(file_name, torch.float, deform_dict, mask, device)
+    if setups['flip']: 
+        Idef = torch.flip(Idef, [0]) 
+    return {'bias_field': Idef[None]}
+
+def read_and_deform_surface(exist_keys, task_name, file_name, setups, deform_dict, device, mask, size):
+    Fneg, A, c2 = deform_dict['Fneg'], deform_dict['A'], deform_dict['c2']
+    # NOTE does not support left_hemis for now
+
+    mat = loadmat(file_name.split('.nii')[0] + '.mat')
+
+    Vlw = torch.tensor(mat['Vlw'], dtype=torch.float, device=device)
+    Flw = torch.tensor(mat['Flw'], dtype=torch.int, device=device)
+    Vrw = torch.tensor(mat['Vrw'], dtype=torch.float, device=device)
+    Frw = torch.tensor(mat['Frw'], dtype=torch.int, device=device)
+    Vlp = torch.tensor(mat['Vlp'], dtype=torch.float, device=device)
+    Flp = torch.tensor(mat['Flp'], dtype=torch.int, device=device)
+    Vrp = torch.tensor(mat['Vrp'], dtype=torch.float, device=device)
+    Frp = torch.tensor(mat['Frp'], dtype=torch.int, device=device)
+
+    Ainv = torch.inverse(A)
+    Vlw -= c2[None, :]
+    Vlw = Vlw @ torch.transpose(Ainv, 0, 1)
+    Vlw += fast_3D_interp_torch(Fneg, Vlw[:, 0] + c2[0], Vlw[:, 1]+c2[1], Vlw[:, 2] + c2[2])
+    Vlw += c2[None, :]
+    Vrw -= c2[None, :]
+    Vrw = Vrw @ torch.transpose(Ainv, 0, 1)
+    Vrw += fast_3D_interp_torch(Fneg, Vrw[:, 0] + c2[0], Vrw[:, 1]+c2[1], Vrw[:, 2] + c2[2])
+    Vrw += c2[None, :]
+    Vlp -= c2[None, :]
+    Vlp = Vlp @ torch.transpose(Ainv, 0, 1)
+    Vlp += fast_3D_interp_torch(Fneg, Vlp[:, 0] + c2[0], Vlp[:, 1] + c2[1], Vlp[:, 2] + c2[2])
+    Vlp += c2[None, :]
+    Vrp -= c2[None, :]
+    Vrp = Vrp @ torch.transpose(Ainv, 0, 1)
+    Vrp += fast_3D_interp_torch(Fneg, Vrp[:, 0] + c2[0], Vrp[:, 1] + c2[1], Vrp[:, 2] + c2[2])
+    Vrp += c2[None, :]
+
+    if setups['flip']: 
+        Vlw[:, 0] = size[0] - 1 - Vlw[:, 0]
+        Vrw[:, 0] = size[0] - 1 - Vrw[:, 0]
+        Vlp[:, 0] = size[0] - 1 - Vlp[:, 0]
+        Vrp[:, 0] = size[0] - 1 - Vrp[:, 0]
+        Vlw, Vrw = Vrw, Vlw
+        Vlp, Vrp = Vrp, Vlp
+        Flw, Frw = Frw, Flw
+        Flp, Frp = Frp, Flp
+
+    print(Vlw.shape) # 131148
+    print(Vlp.shape) # 131148
+
+    print(Vrw.shape) # 131720
+    print(Vrp.shape) # 131720
+
+    print(Flw.shape) # 262292
+    print(Flp.shape) # 262292
+
+    print(Frw.shape) # 263436
+    print(Frp.shape) # 263436
+    #return torch.stack([Vlw, Flw, Vrw, Frw, Vlp, Flp, Vrp, Frp])
+    return {'Vlw': Vlw, 'Flw': Flw, 'Vrw': Vrw, 'Frw': Frw, 'Vlp': Vlp, 'Flp': Flp, 'Vrp': Vrp, 'Frp': Frp}
+    
+
+#####################################
+#########  Pathology Shape  #########
+#####################################
+
+
+def augment_pathology(Pprob, pde_func, t, shape_gen_args, device):
+    Pprob = torch.squeeze(Pprob) 
+
+    nt = np.random.randint(1, shape_gen_args.max_nt+1)
+    if nt <= 1:  
+        return Pprob
+
+    pde_func.V_dict = generate_velocity_3d(Pprob.shape, shape_gen_args.perlin_res, shape_gen_args.V_multiplier, device)
+
+    #start_time = time.time()
+    Pprob = odeint(pde_func, Pprob[None], t[:nt], 
+                    shape_gen_args.dt, 
+                    method = shape_gen_args.integ_method)[-1, 0] # (last_t, n_batch=1, s, r, c)
+    # total_time = time.time() - start_time
+    #total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    #print('Time {} for {} time points'.format(total_time_str, nt)) 
+
+    
+    return Pprob
+
+
+#####################################
+######### Augmentation Func #########
+#####################################
+    
+
+def add_gamma_transform(I, aux_dict, cfg, device, **kwargs):
+    gamma = torch.tensor(np.exp(cfg.gamma_std * np.random.randn(1)[0]), dtype=float, device=device)
+    I_gamma = 300.0 * (I / 300.0) ** gamma
+    #aux_dict.update({'gamma': gamma}) # uncomment if you want to save gamma for later use
+    return I_gamma, aux_dict
+
+def add_bias_field(I, aux_dict, cfg, input_mode, setups, size, device, **kwargs):
+    if input_mode == 'CT':
+        aux_dict.update({'high_res': I})
+        return I, aux_dict
+    
+    bf_scale = cfg.bf_scale_min + np.random.rand(1) * (cfg.bf_scale_max - cfg.bf_scale_min)
+    size_BF_small = np.round(bf_scale * np.array(size)).astype(int).tolist()
+    if setups['photo_mode']:
+        size_BF_small[1] = np.round(size[1]/setups['spac']).astype(int)
+    BFsmall = torch.tensor(cfg.bf_std_min + (cfg.bf_std_max - cfg.bf_std_min) * np.random.rand(1), dtype=torch.float, device=device) * \
+        torch.randn(size_BF_small, dtype=torch.float, device=device)
+    BFlog = myzoom_torch(BFsmall, np.array(size) / size_BF_small)
+    BF = torch.exp(BFlog)
+    I_bf = I * BF
+    aux_dict.update({'BFlog': BFlog, 'high_res': I_bf})
+    return I_bf, aux_dict
+
+def resample_resolution(I, aux_dict, setups, res, size, device, **kwargs):
+    stds = (0.85 + 0.3 * np.random.rand()) * np.log(5) /np.pi * setups['thickness'] / res
+    stds[setups['thickness']<=res] = 0.0 # no blur if thickness is equal to the resolution of the training data
+    I_blur = gaussian_blur_3d(I, stds, device)
+    new_size = (np.array(size) * res / setups['resolution']).astype(int)
+
+    factors = np.array(new_size) / np.array(size)
+    delta = (1.0 - factors) / (2.0 * factors)
+    vx = np.arange(delta[0], delta[0] + new_size[0] / factors[0], 1 / factors[0])[:new_size[0]]
+    vy = np.arange(delta[1], delta[1] + new_size[1] / factors[1], 1 / factors[1])[:new_size[1]]
+    vz = np.arange(delta[2], delta[2] + new_size[2] / factors[2], 1 / factors[2])[:new_size[2]]
+    II, JJ, KK = np.meshgrid(vx, vy, vz, sparse=False, indexing='ij')
+    II = torch.tensor(II, dtype=torch.float, device=device)
+    JJ = torch.tensor(JJ, dtype=torch.float, device=device)
+    KK = torch.tensor(KK, dtype=torch.float, device=device)
+
+    I_small = fast_3D_interp_torch(I_blur, II, JJ, KK) 
+    aux_dict.update({'factors': factors})
+    return I_small, aux_dict
+
+
+def resample_resolution_photo(I, aux_dict, setups, res, size, device, **kwargs):
+    stds = (0.85 + 0.3 * np.random.rand()) * np.log(5) /np.pi * setups['thickness'] / res
+    stds[setups['thickness']<=res] = 0.0 # no blur if thickness is equal to the resolution of the training data
+    I_blur = gaussian_blur_3d(I, stds, device)
+    new_size = (np.array(size) * res / setups['resolution']).astype(int)
+
+    factors = np.array(new_size) / np.array(size)
+    delta = (1.0 - factors) / (2.0 * factors)
+    vx = np.arange(delta[0], delta[0] + new_size[0] / factors[0], 1 / factors[0])[:new_size[0]]
+    vy = np.arange(delta[1], delta[1] + new_size[1] / factors[1], 1 / factors[1])[:new_size[1]]
+    vz = np.arange(delta[2], delta[2] + new_size[2] / factors[2], 1 / factors[2])[:new_size[2]]
+    II, JJ, KK = np.meshgrid(vx, vy, vz, sparse=False, indexing='ij')
+    II = torch.tensor(II, dtype=torch.float, device=device)
+    JJ = torch.tensor(JJ, dtype=torch.float, device=device)
+    KK = torch.tensor(KK, dtype=torch.float, device=device)
+
+    I_small = fast_3D_interp_torch(I_blur, II, JJ, KK) 
+    aux_dict.update({'factors': factors})
+    return I_small, aux_dict
+
+
+def add_noise(I, aux_dict, cfg, device, **kwargs):
+    noise_std = torch.tensor(cfg.noise_std_min + (cfg.noise_std_max - cfg.noise_std_min) * np.random.rand(1), dtype=torch.float, device=device)
+    I_noisy = I + noise_std * torch.randn(I.shape, dtype=torch.float, device=device)
+    I_noisy[I_noisy < 0] = 0
+    #aux_dict.update({'noise_std': noise_std}) # uncomment if you want to save noise_std for later use
+    return I_noisy, aux_dict
+    
+    
+#####################################
+#####################################
+
+
+# map SynthSeg right to left labels for contrast synthesis
+right_to_left_dict = {
+    41: 2,
+    42: 3,
+    43: 4,
+    44: 5,
+    46: 7,
+    47: 8,
+    49: 10,
+    50: 11,
+    51: 12,
+    52: 13,
+    53: 17,
+    54: 18,
+    58: 26,
+    60: 28
+}
+
+# based on merged left & right SynthSeg labels
+ct_brightness_group = {
+    'darker': [4, 5, 14, 15, 24, 31, 72], # ventricles, CSF
+    'dark': [2, 7, 16, 77, 30], # white matter
+    'bright': [3, 8, 17, 18, 28, 10, 11, 12, 13, 26], # grey matter (cortex, hippocampus, amggdala, ventral DC), thalamus, ganglia (nucleus (putamen, pallidus, accumbens), caudate)
+    'brighter': [], # skull, pineal gland, choroid plexus
+}

README.md

@@ -1,3 +1,91 @@
----
-license: apache-2.0
----
+
+## <p align="center">[A Modality-agnostic Multi-task Foundation Model for Human Brain Imaging](https://arxiv.org/abs/2509.00549)</p>
+
+**<p align="center">Peirong Liu<sup>1,2</sup>, Oula Puonti<sup>2</sup>, Xiaoling Hu<sup>2</sup>, Karthik Gopinath<sup>2</sup>, Annabel Sorby-Adams<sup>2</sup>, Daniel C. Alexander<sup>3</sup>, Juan Eugenio Iglesias<sup>2,3,4</sup></p>**
+
+ 
+<p align="center">
+<sup>1</sup>Johns Hopkins University<br />
+<sup>2</sup>Harvard Medical School and Massachusetts General Hospital<br />
+<sup>3</sup>University College London <br />
+<sup>4</sup>Massachusetts Institute of Technology
+</p>
+
+<p align="center">
+  <img src="./assets/overview.png" alt="drawing", width="650"/>
+</p>
+
+
+This is the official repository for our preprint: A Modality-agnostic Multi-task Foundation Model for Human Brain Imaging [[arXiv]](https://arxiv.org/abs/2509.00549)<br />
+More detailed and organized instructions are coming soon...
+
+## Environment
+Training and evaluation environment: Python 3.11.4, PyTorch 2.0.1, CUDA 12.2. Run the following command to install required packages.
+```
+conda create -n pre python=3.11
+conda activate pre
+
+git clone https://github.com/jhuldr/BrainFM
+cd /path/to/brainfm
+pip install -r requirements.txt 
+```
+
+
+## Generator
+```
+cd scripts 
+python demo_generator.py
+```
+
+### Generator setups
+Setups are in cfgs/generator, default setups are in default.yaml. A customized setup example can be found in train/brain_id.yaml, where several Brain-ID-specific setups are added. During Config reading/implementation, customized yaml will overwrite default.yaml if they have the same keys.
+
+dataset_setups: information for all datasets, in Generator/constants.py<br>
+augmentation_funcs: augmentation functions and steps, in Generator/constants.py<br>
+processing_funcs: image processing functions for each modality/task, in Generator/constants.py<br>
+
+dataset_names: dataset name list, paths setups in Generator/constants.py<br>
+mix_synth_prob: if the input mode is synthesizing, then probability for blending synth with real images<br>
+dataset_option: generator types, could be BaseGen or customized generator<br>
+task: switch on/off individual training tasks 
+
+### Base generator module
+```
+cd Generator
+python datasets.py
+```
+The dataset paths setups are in constants.py. In datasets.py, different datasets been used are fomulated as a list of dataset names.
+
+A customized data generator module example can be found in datasets.py -- BrainIDGen. 
+
+
+Refer to "__getitem__" function. Specifically, it includes: <br>
+(1) read original input: could be either generation labels or real images;<br>
+(2) generate augmentation setups and deformation fields; <br>
+(3) read target(s) according to the assigned tasks -- here I seperate the processing functions for each item/modality, in case we want different processing steps for them; <br>
+(4) augment input sample: either synthesized or real image input.
+
+
+
+(Some of the functions are leaved blank for now.)
+
+
+
+## Trainer
+```
+cd scripts 
+python train.py
+```
+
+## Downloads
+The pre-trained model weight is available on [OneDrive](https://livejohnshopkins-my.sharepoint.com/:u:/g/personal/pliu53_jh_edu/EZ_BJ7K6pMJEj9hZ8SA51GYBxH_Nan4fA3a-s4udwvVRog?e=nwZ7JC).
+
+
+## Citation
+```bibtex
+@article{Liu_2025_BrainFM,
+    author    = {Liu, Peirong and Puonti, Oula and Hu, Xiaoling and Gopinath, Karthik and Sorby-Adams, Annabel and Alexander, Daniel C. and Iglesias, Juan E.},
+    title     = {A Modality-agnostic Multi-task Foundation Model for Human Brain Imaging},
+    booktitle = {arXiv preprint arXiv:2509.00549},
+    year      = {2025},
+}

ShapeID/DiffEqs/FD.py

@@ -0,0 +1,525 @@
+"""
+*finite_difference.py* is the main package to compute finite differences in 
+1D, 2D, and 3D on numpy arrays (class FD_np) and pytorch tensors (class FD_torch).
+The package supports first and second order derivatives and Neumann and linear extrapolation
+boundary conditions (though the latter have not been tested extensively yet).
+"""
+from __future__ import absolute_import
+
+# from builtins import object
+from abc import ABCMeta, abstractmethod
+
+import torch
+from torch.autograd import Variable
+import numpy as np
+from future.utils import with_metaclass
+
+class FD(with_metaclass(ABCMeta, object)):
+    """
+    *FD* is the abstract class for finite differences. It includes most of the actual finite difference code, 
+    but requires the definition (in a derived class) of the methods *get_dimension*, *create_zero_array*, and *get_size_of_array*.
+    In this way the numpy and pytorch versions can easily be derived. All the method expect BxXxYxZ format (i.e., they process a batch at a time)
+    """
+
+    def __init__(self, spacing, bcNeumannZero=True):
+        """
+        Constructor        
+        :param spacing: 1D numpy array defining the spatial spacing, e.g., [0.1,0.1,0.1] for a 3D image  
+        :param bcNeumannZero: Defines the boundary condition. If set to *True* (default) zero Neumann boundary conditions
+            are imposed. If set to *False* linear extrapolation is used (this is still experimental, but may be beneficial 
+            for better boundary behavior)
+        """
+
+        self.dim = len(spacing) # In my code, data_spacing is a list # spacing.size
+        """spatial dimension"""
+        self.spacing = np.ones(self.dim)
+        """spacing"""
+        self.bcNeumannZero = bcNeumannZero # if false then linear interpolation
+        """should Neumann boundary conditions be used? (otherwise linear extrapolation)"""
+        if len(spacing) == 1: #spacing.size==1:
+            self.spacing[0] = spacing[0]
+        elif len(spacing) == 2: # spacing.size==2:
+            self.spacing[0] = spacing[0]
+            self.spacing[1] = spacing[1]
+        elif len(spacing) == 3: # spacing.size==3:
+            self.spacing[0] = spacing[0]
+            self.spacing[1] = spacing[1]
+            self.spacing[2] = spacing[2]
+        else:
+            print('Current dimension:', len(spacing))
+            raise ValueError('Finite differences are only supported in dimensions 1 to 3')
+
+    def dXb(self,I):
+        """
+        Backward difference in x direction:
+        :math:`\\frac{dI(i)}{dx}\\approx\\frac{I_i-I_{i-1}}{h_x}`
+        :param I: Input image  
+        :return: Returns the first derivative in x direction using backward differences
+        """
+        return (I-self.xm(I))/self.spacing[0]
+
+    def dXf(self,I):
+        """
+        Forward difference in x direction:
+        :math:`\\frac{dI(i)}{dx}\\approx\\frac{I_{i+1}-I_{i}}{h_x}`
+        
+        :param I: Input image
+        :return: Returns the first derivative in x direction using forward differences
+        """
+        return (self.xp(I)-I)/self.spacing[0]
+
+    def dXc(self,I):
+        """
+        Central difference in x direction:
+        :math:`\\frac{dI(i)}{dx}\\approx\\frac{I_{i+1}-I_{i-1}}{2h_x}`
+        
+        :param I: Input image
+        :return: Returns the first derivative in x direction using central differences
+        """
+        return (self.xp(I)-self.xm(I))/(2*self.spacing[0])
+
+    def ddXc(self,I):
+        """
+        Second deriative in x direction
+        
+        :param I: Input image
+        :return: Returns the second derivative in x direction
+        """
+        return (self.xp(I)-2*I+self.xm(I))/(self.spacing[0]**2)
+
+    def dYb(self,I):
+        """
+        Same as dXb, but for the y direction
+        
+        :param I: Input image
+        :return: Returns the first derivative in y direction using backward differences
+        """
+        return (I-self.ym(I))/self.spacing[1]
+
+    def dYf(self,I):
+        """
+        Same as dXf, but for the y direction
+        
+        :param I: Input image
+        :return: Returns the first derivative in y direction using forward differences
+        """
+        return (self.yp(I)-I)/self.spacing[1]
+
+    def dYc(self,I):
+        """
+        Same as dXc, but for the y direction
+        
+        :param I: Input image
+        :return: Returns the first derivative in y direction using central differences
+        """
+        return (self.yp(I)-self.ym(I))/(2*self.spacing[1])
+
+    def ddYc(self,I):
+        """
+        Same as ddXc, but for the y direction
+        
+        :param I: Input image
+        :return: Returns the second derivative in the y direction
+        """
+        return (self.yp(I)-2*I+self.ym(I))/(self.spacing[1]**2)
+
+    def dZb(self,I):
+        """
+        Same as dXb, but for the z direction
+        
+        :param I: Input image 
+        :return: Returns the first derivative in the z direction using backward differences
+        """
+        return (I - self.zm(I))/self.spacing[2]
+
+    def dZf(self, I):
+        """
+        Same as dXf, but for the z direction
+        
+        :param I: Input image
+        :return: Returns the first derivative in the z direction using forward differences
+        """
+        return (self.zp(I)-I)/self.spacing[2]
+
+    def dZc(self, I):
+        """
+        Same as dXc, but for the z direction
+        
+        :param I: Input image
+        :return: Returns the first derivative in the z direction using central differences
+        """
+        return (self.zp(I)-self.zm(I))/(2*self.spacing[2])
+
+    def ddZc(self,I):
+        """
+        Same as ddXc, but for the z direction
+        
+        :param I: Input iamge
+        :return: Returns the second derivative in the z direction 
+        """
+        return (self.zp(I)-2*I+self.zm(I))/(self.spacing[2]**2)
+
+    def lap(self, I):
+        """
+        Compute the Lapacian of an image
+        !!!!!!!!!!!
+        IMPORTANT:
+        ALL THE FOLLOWING IMPLEMENTED CODE ADD 1 ON DIMENSION, WHICH REPRESENT BATCH DIMENSION.
+        THIS IS FOR COMPUTATIONAL EFFICIENCY.
+        
+        :param I: Input image [batch, X,Y,Z]
+        :return: Returns the Laplacian
+        """
+        ndim = self.getdimension(I)
+        if ndim == 1+1:
+            return self.ddXc(I)
+        elif ndim == 2+1:
+            return (self.ddXc(I) + self.ddYc(I))
+        elif ndim == 3+1:
+            return (self.ddXc(I) + self.ddYc(I) + self.ddZc(I))
+        else:
+            raise ValueError('Finite differences are only supported in dimensions 1 to 3')
+
+    def grad_norm_sqr_c(self, I):
+        """
+        Computes the gradient norm of an image
+        !!!!!!!!!!!
+        IMPORTANT:
+        ALL THE FOLLOWING IMPLEMENTED CODE ADD 1 ON DIMENSION, WHICH REPRESENT BATCH DIMENSION.
+        THIS IS FOR COMPUTATIONAL EFFICIENCY.
+        :param I: Input image [batch, X,Y,Z]
+        :return: returns ||grad I||^2
+        """
+        ndim = self.getdimension(I)
+        if ndim == 1 + 1:
+            return self.dXc(I)**2
+        elif ndim == 2 + 1:
+            return (self.dXc(I)**2 + self.dYc(I)**2)
+        elif ndim == 3 + 1:
+            return (self.dXc(I)**2 + self.dYc(I)**2 + self.dZc(I)**2)
+        else:
+            raise ValueError('Finite differences are only supported in dimensions 1 to 3')
+
+    def grad_norm_sqr_f(self, I):
+        """
+        Computes the gradient norm of an image
+        !!!!!!!!!!!
+        IMPORTANT:
+        ALL THE FOLLOWING IMPLEMENTED CODE ADD 1 ON DIMENSION, WHICH REPRESENT BATCH DIMENSION.
+        THIS IS FOR COMPUTATIONAL EFFICIENCY.
+        :param I: Input image [batch, X,Y,Z]
+        :return: returns ||grad I||^2
+        """
+        ndim = self.getdimension(I)
+        if ndim == 1 + 1:
+            return self.dXf(I)**2
+        elif ndim == 2 + 1:
+            return (self.dXf(I)**2 + self.dYf(I)**2)
+        elif ndim == 3 + 1:
+            return (self.dXf(I)**2 + self.dYf(I)**2 + self.dZf(I)**2)
+        else:
+            raise ValueError('Finite differences are only supported in dimensions 1 to 3')
+
+    def grad_norm_sqr_b(self, I):
+        """
+        Computes the gradient norm of an image
+        !!!!!!!!!!!
+        IMPORTANT:
+        ALL THE FOLLOWING IMPLEMENTED CODE ADD 1 ON DIMENSION, WHICH REPRESENT BATCH DIMENSION.
+        THIS IS FOR COMPUTATIONAL EFFICIENCY.
+        :param I: Input image [batch, X,Y,Z]
+        :return: returns ||grad I||^2
+        """
+        ndim = self.getdimension(I)
+        if ndim == 1 + 1:
+            return self.dXb(I)**2
+        elif ndim == 2 + 1:
+            return (self.dXb(I)**2 + self.dYb(I)**2)
+        elif ndim == 3 + 1:
+            return (self.dXb(I)**2 + self.dYb(I)**2 + self.dZb(I)**2)
+        else:
+            raise ValueError('Finite differences are only supported in dimensions 1 to 3')
+
+    @abstractmethod
+    def getdimension(self,I):
+        """
+        Abstract method to return the dimension of an input image I
+        
+        :param I: Input image 
+        :return: Returns the dimension of the image I
+        """
+        pass
+
+    @abstractmethod
+    def create_zero_array(self, sz):
+        """
+        Abstract method to create a zero array of a given size, sz. E.g., sz=[10,2,5]
+        
+        :param sz: Size array
+        :return: Returns a zero array of the specified size
+        """
+        pass
+
+    @abstractmethod
+    def get_size_of_array(self, A):
+        """
+        Abstract method to return the size of an array (as a vector)
+        
+        :param A: Input array
+        :return: Returns its size (e.g., [5,10] or [3,4,6]
+        """
+        pass
+
+    def xp(self,I):
+        """
+        !!!!!!!!!!!
+        IMPORTANT:
+        ALL THE FOLLOWING IMPLEMENTED CODE ADD 1 ON DIMENSION, WHICH REPRESENT BATCH DIMENSION.
+        THIS IS FOR COMPUTATIONAL EFFICIENCY.
+        Returns the values for x-index incremented by one (to the right in 1D)
+        
+        :param I: Input image [batch, X, Y,Z]
+        :return: Image with values at an x-index one larger
+        """
+        rxp = self.create_zero_array( self.get_size_of_array( I ) )
+        ndim = self.getdimension(I)
+        if ndim == 1+1:
+            rxp[:,0:-1] = I[:,1:]
+            if self.bcNeumannZero:
+                rxp[:,-1] = I[:,-1]
+            else:
+                rxp[:,-1] = 2*I[:,-1]-I[:,-2]
+        elif ndim == 2+1:
+            rxp[:,0:-1,:] = I[:,1:,:]
+            if self.bcNeumannZero:
+                rxp[:,-1,:] = I[:,-1,:]
+            else:
+                rxp[:,-1,:] = 2*I[:,-1,:]-I[:,-2,:]
+        elif ndim == 3+1:
+            rxp[:,0:-1,:,:] = I[:,1:,:,:]
+            if self.bcNeumannZero:
+                rxp[:,-1,:,:] = I[:,-1,:,:]
+            else:
+                rxp[:,-1,:,:] = 2*I[:,-1,:,:]-I[:,-2,:,:]
+        else:
+            raise ValueError('Finite differences are only supported in dimensions 1 to 3')
+        return rxp
+
+    def xm(self,I):
+        """
+        !!!!!!!!!!!
+        IMPORTANT:
+        ALL THE FOLLOWING IMPLEMENTED CODE ADD 1 ON DIMENSION, WHICH REPRESENT BATCH DIMENSION.
+        THIS IS FOR COMPUTATIONAL EFFICIENCY.
+        Returns the values for x-index decremented by one (to the left in 1D)
+        
+        :param I: Input image [batch, X, Y, Z]
+        :return: Image with values at an x-index one smaller
+        """
+        rxm = self.create_zero_array( self.get_size_of_array( I ) )
+        ndim = self.getdimension(I)
+        if ndim == 1+1:
+            rxm[:,1:] = I[:,0:-1]
+            if self.bcNeumannZero:
+                rxm[:,0] = I[:,0]
+            else:
+                rxm[:,0] = 2*I[:,0]-I[:,1]
+        elif ndim == 2+1:
+            rxm[:,1:,:] = I[:,0:-1,:]
+            if self.bcNeumannZero:
+                rxm[:,0,:] = I[:,0,:]
+            else:
+                rxm[:,0,:] = 2*I[:,0,:]-I[:,1,:]
+        elif ndim == 3+1:
+            rxm[:,1:,:,:] = I[:,0:-1,:,:]
+            if self.bcNeumannZero:
+                rxm[:,0,:,:] = I[:,0,:,:]
+            else:
+                rxm[:,0,:,:] = 2*I[:,0,:,:]-I[:,1,:,:]
+        else:
+            raise ValueError('Finite differences are only supported in dimensions 1 to 3')
+        return rxm
+
+    def yp(self, I):
+        """
+        !!!!!!!!!!!
+        IMPORTANT:
+        ALL THE FOLLOWING IMPLEMENTED CODE ADD 1 ON DIMENSION, WHICH REPRESENT BATCH DIMENSION.
+        THIS IS FOR COMPUTATIONAL EFFICIENCY.
+        Same as xp, but for the y direction
+        
+        :param I: Input image
+        :return: Image with values at y-index one larger
+        """
+        ryp = self.create_zero_array( self.get_size_of_array( I ) )
+        ndim = self.getdimension(I)
+        if ndim == 2+1:
+            ryp[:,:,0:-1] = I[:,:,1:]
+            if self.bcNeumannZero:
+                ryp[:,:,-1] = I[:,:,-1]
+            else:
+                ryp[:,:,-1] = 2*I[:,:,-1]-I[:,:,-2]
+        elif ndim == 3+1:
+            ryp[:,:,0:-1,:] = I[:,:,1:,:]
+            if self.bcNeumannZero:
+                ryp[:,:,-1,:] = I[:,:,-1,:]
+            else:
+                ryp[:,:,-1,:] = 2*I[:,:,-1,:]-I[:,:,-2,:]
+        else:
+            print('Current dimension:', ndim-1)
+            raise ValueError('Finite differences are only supported in dimensions 1 to 3')
+        return ryp
+
+    def ym(self, I):
+        """
+        Same as xm, but for the y direction
+        !!!!!!!!!!!
+        IMPORTANT:
+        ALL THE FOLLOWING IMPLEMENTED CODE ADD 1 ON DIMENSION, WHICH REPRESENT BATCH DIMENSION.
+        THIS IS FOR COMPUTATIONAL EFFICIENCY.
+        Returns the values for x-index decremented by one (to the left in 1D)
+        :param I: Input image [batch, X, Y, Z]
+        :return: Image with values at y-index one smaller
+        """
+        rym = self.create_zero_array( self.get_size_of_array( I ) )
+        ndim = self.getdimension(I)
+        if ndim == 2+1:
+            rym[:,:,1:] = I[:,:,0:-1]
+            if self.bcNeumannZero:
+                rym[:,:,0] = I[:,:,0]
+            else:
+                rym[:,:,0] = 2*I[:,:,0]-I[:,:,1]
+        elif ndim == 3+1:
+            rym[:,:,1:,:] = I[:,:,0:-1,:]
+            if self.bcNeumannZero:
+                rym[:,:,0,:] = I[:,:,0,:]
+            else:
+                rym[:,:,0,:] = 2*I[:,:,0,:]-I[:,:,1,:]
+        else:
+            raise ValueError('Finite differences are only supported in dimensions 1 to 3')
+        return rym
+
+    def zp(self, I):
+        """
+        Same as xp, but for the z direction
+        
+        !!!!!!!!!!!
+        IMPORTANT:
+        ALL THE FOLLOWING IMPLEMENTED CODE ADD 1 ON DIMENSION, WHICH REPRESENT BATCH DIMENSION.
+        THIS IS FOR COMPUTATIONAL EFFICIENCY.
+        Returns the values for x-index decremented by one (to the left in 1D)
+        :param I: Input image [batch, X, Y, Z]
+        :return: Image with values at z-index one larger
+        """
+        rzp = self.create_zero_array( self.get_size_of_array( I ) )
+        ndim = self.getdimension(I)
+        if ndim == 3+1:
+            rzp[:,:,:,0:-1] = I[:,:,:,1:]
+            if self.bcNeumannZero:
+                rzp[:,:,:,-1] = I[:,:,:,-1]
+            else:
+                rzp[:,:,:,-1] = 2*I[:,:,:,-1]-I[:,:,:,-2]
+        else:
+            raise ValueError('Finite differences are only supported in dimensions 1 to 3')
+        return rzp
+
+    def zm(self, I):
+        """
+        Same as xm, but for the z direction
+        
+        !!!!!!!!!!!
+        IMPORTANT:
+        ALL THE FOLLOWING IMPLEMENTED CODE ADD 1 ON DIMENSION, WHICH REPRESENT BATCH DIMENSION.
+        THIS IS FOR COMPUTATIONAL EFFICIENCY.
+        Returns the values for x-index decremented by one (to the left in 1D)
+        :param I: Input image [batch, X, Y, Z]
+        :return: Image with values at z-index one smaller
+        """
+        rzm = self.create_zero_array( self.get_size_of_array( I ) )
+        ndim = self.getdimension(I)
+        if ndim == 3+1:
+            rzm[:,:,:,1:] = I[:,:,:,0:-1]
+            if self.bcNeumannZero:
+                rzm[:,:,:,0] = I[:,:,:,0]
+            else:
+                rzm[:,:,:,0] = 2*I[:,:,:,0]-I[:,:,:,1]
+        else:
+            raise ValueError('Finite differences are only supported in dimensions 1 to 3')
+        return rzm
+
+
+class FD_np(FD):
+    """
+    Defnitions of the abstract methods for numpy
+    """
+
+    def __init__(self,spacing,bcNeumannZero=True):
+        """
+        Constructor for numpy finite differences
+        :param spacing: spatial spacing (array with as many entries as there are spatial dimensions)
+        :param bcNeumannZero: Specifies if zero Neumann conditions should be used (if not, uses linear extrapolation)
+        """
+        super(FD_np, self).__init__(spacing,bcNeumannZero)
+
+    def getdimension(self,I):
+        """
+        Returns the dimension of an image
+        :param I: input image
+        :return: dimension of the input image
+        """
+        return I.ndim
+
+    def create_zero_array(self, sz):
+        """
+        Creates a zero array
+        :param sz: size of the zero array, e.g., [3,4,2]
+        :return: the zero array
+        """
+        return np.zeros( sz )
+
+    def get_size_of_array(self, A):
+        """
+        Returns the size (shape in numpy) of an array
+        :param A: input array
+        :return: shape/size
+        """
+        return A.shape
+
+
+class FD_torch(FD):
+    """
+    Defnitions of the abstract methods for torch
+    """
+
+    def __init__(self,spacing,device,bcNeumannZero=True):
+        """
+          Constructor for torch finite differences
+          :param spacing: spatial spacing (array with as many entries as there are spatial dimensions)
+          :param bcNeumannZero: Specifies if zero Neumann conditions should be used (if not, uses linear extrapolation)
+          """
+        super(FD_torch, self).__init__(spacing,bcNeumannZero)
+        self.device = device
+
+    def getdimension(self,I):
+        """
+        Returns the dimension of an image
+        :param I: input image
+        :return: dimension of the input image 
+        """
+        return I.dim()
+
+    def create_zero_array(self, sz):
+        """
+        Creats a zero array
+        :param sz: size of the array, e.g., [3,4,2]
+        :return: the zero array
+        """
+        return  torch.zeros(sz).float().to(self.device)
+
+    def get_size_of_array(self, A):
+        """
+        Returns the size (size()) of an array
+        :param A: input array
+        :return: shape/size
+        """
+        return A.size()

ShapeID/DiffEqs/adams.py

@@ -0,0 +1,170 @@
+import collections
+import torch
+from ShapeID.DiffEqs.solvers import AdaptiveStepsizeODESolver
+from ShapeID.DiffEqs.misc import (
+    _handle_unused_kwargs, _select_initial_step, _convert_to_tensor, _scaled_dot_product, _is_iterable,
+    _optimal_step_size, _compute_error_ratio
+)
+
+_MIN_ORDER = 1
+_MAX_ORDER = 12
+
+gamma_star = [
+    1, -1 / 2, -1 / 12, -1 / 24, -19 / 720, -3 / 160, -863 / 60480, -275 / 24192, -33953 / 3628800, -0.00789255,
+    -0.00678585, -0.00592406, -0.00523669, -0.0046775, -0.00421495, -0.0038269
+]
+
+
+class _VCABMState(collections.namedtuple('_VCABMState', 'y_n, prev_f, prev_t, next_t, phi, order')):
+    """Saved state of the variable step size Adams-Bashforth-Moulton solver as described in
+
+        Solving Ordinary Differential Equations I - Nonstiff Problems III.5
+        by Ernst Hairer, Gerhard Wanner, and Syvert P Norsett.
+    """
+
+
+def g_and_explicit_phi(prev_t, next_t, implicit_phi, k):
+    curr_t = prev_t[0]
+    dt = next_t - prev_t[0]
+
+    g = torch.empty(k + 1).to(prev_t[0])
+    explicit_phi = collections.deque(maxlen=k)
+    beta = torch.tensor(1).to(prev_t[0])
+
+    g[0] = 1
+    c = 1 / torch.arange(1, k + 2).to(prev_t[0])
+    explicit_phi.append(implicit_phi[0])
+
+    for j in range(1, k):
+        beta = (next_t - prev_t[j - 1]) / (curr_t - prev_t[j]) * beta
+        beat_cast = beta.to(implicit_phi[j][0])
+        explicit_phi.append(tuple(iphi_ * beat_cast for iphi_ in implicit_phi[j]))
+
+        c = c[:-1] - c[1:] if j == 1 else c[:-1] - c[1:] * dt / (next_t - prev_t[j - 1])
+        g[j] = c[0]
+
+    c = c[:-1] - c[1:] * dt / (next_t - prev_t[k - 1])
+    g[k] = c[0]
+
+    return g, explicit_phi
+
+
+def compute_implicit_phi(explicit_phi, f_n, k):
+    k = min(len(explicit_phi) + 1, k)
+    implicit_phi = collections.deque(maxlen=k)
+    implicit_phi.append(f_n)
+    for j in range(1, k):
+        implicit_phi.append(tuple(iphi_ - ephi_ for iphi_, ephi_ in zip(implicit_phi[j - 1], explicit_phi[j - 1])))
+    return implicit_phi
+
+
+class VariableCoefficientAdamsBashforth(AdaptiveStepsizeODESolver):
+
+    def __init__(
+        self, func, y0, rtol, atol, implicit=True, max_order=_MAX_ORDER, safety=0.9, ifactor=10.0, dfactor=0.2,
+        **unused_kwargs
+    ):
+        _handle_unused_kwargs(self, unused_kwargs)
+        del unused_kwargs
+
+        self.func = func
+        self.y0 = y0
+        self.rtol = rtol if _is_iterable(rtol) else [rtol] * len(y0)
+        self.atol = atol if _is_iterable(atol) else [atol] * len(y0)
+        self.implicit = implicit
+        self.max_order = int(max(_MIN_ORDER, min(max_order, _MAX_ORDER)))
+        self.safety = _convert_to_tensor(safety, dtype=torch.float64, device=y0[0].device)
+        self.ifactor = _convert_to_tensor(ifactor, dtype=torch.float64, device=y0[0].device)
+        self.dfactor = _convert_to_tensor(dfactor, dtype=torch.float64, device=y0[0].device)
+
+    def before_integrate(self, t):
+        prev_f = collections.deque(maxlen=self.max_order + 1)
+        prev_t = collections.deque(maxlen=self.max_order + 1)
+        phi = collections.deque(maxlen=self.max_order)
+
+        t0 = t[0]
+        f0 = self.func(t0.type_as(self.y0[0]), self.y0)
+        prev_t.appendleft(t0)
+        prev_f.appendleft(f0)
+        phi.appendleft(f0)
+        first_step = _select_initial_step(self.func, t[0], self.y0, 2, self.rtol[0], self.atol[0], f0=f0).to(t)
+
+        self.vcabm_state = _VCABMState(self.y0, prev_f, prev_t, next_t=t[0] + first_step, phi=phi, order=1)
+
+    def advance(self, final_t):
+        final_t = _convert_to_tensor(final_t).to(self.vcabm_state.prev_t[0])
+        while final_t > self.vcabm_state.prev_t[0]:
+            self.vcabm_state = self._adaptive_adams_step(self.vcabm_state, final_t)
+        assert final_t == self.vcabm_state.prev_t[0]
+        return self.vcabm_state.y_n
+
+    def _adaptive_adams_step(self, vcabm_state, final_t):
+        y0, prev_f, prev_t, next_t, prev_phi, order = vcabm_state
+        if next_t > final_t:
+            next_t = final_t
+        dt = (next_t - prev_t[0])
+        dt_cast = dt.to(y0[0])
+
+        # Explicit predictor step.
+        g, phi = g_and_explicit_phi(prev_t, next_t, prev_phi, order)
+        g = g.to(y0[0])
+        p_next = tuple(
+            y0_ + _scaled_dot_product(dt_cast, g[:max(1, order - 1)], phi_[:max(1, order - 1)])
+            for y0_, phi_ in zip(y0, tuple(zip(*phi)))
+        )
+
+        # Update phi to implicit.
+        next_f0 = self.func(next_t.to(p_next[0]), p_next)
+        implicit_phi_p = compute_implicit_phi(phi, next_f0, order + 1)
+
+        # Implicit corrector step.
+        y_next = tuple(
+            p_next_ + dt_cast * g[order - 1] * iphi_ for p_next_, iphi_ in zip(p_next, implicit_phi_p[order - 1])
+        )
+
+        # Error estimation.
+        tolerance = tuple(
+            atol_ + rtol_ * torch.max(torch.abs(y0_), torch.abs(y1_))
+            for atol_, rtol_, y0_, y1_ in zip(self.atol, self.rtol, y0, y_next)
+        )
+        local_error = tuple(dt_cast * (g[order] - g[order - 1]) * iphi_ for iphi_ in implicit_phi_p[order])
+        error_k = _compute_error_ratio(local_error, tolerance)
+        accept_step = (torch.tensor(error_k) <= 1).all()
+
+        if not accept_step:
+            # Retry with adjusted step size if step is rejected.
+            dt_next = _optimal_step_size(dt, error_k, self.safety, self.ifactor, self.dfactor, order=order)
+            return _VCABMState(y0, prev_f, prev_t, prev_t[0] + dt_next, prev_phi, order=order)
+
+        # We accept the step. Evaluate f and update phi.
+        next_f0 = self.func(next_t.to(p_next[0]), y_next)
+        implicit_phi = compute_implicit_phi(phi, next_f0, order + 2)
+
+        next_order = order
+
+        if len(prev_t) <= 4 or order < 3:
+            next_order = min(order + 1, 3, self.max_order)
+        else:
+            error_km1 = _compute_error_ratio(
+                tuple(dt_cast * (g[order - 1] - g[order - 2]) * iphi_ for iphi_ in implicit_phi_p[order - 1]), tolerance
+            )
+            error_km2 = _compute_error_ratio(
+                tuple(dt_cast * (g[order - 2] - g[order - 3]) * iphi_ for iphi_ in implicit_phi_p[order - 2]), tolerance
+            )
+            if min(error_km1 + error_km2) < max(error_k):
+                next_order = order - 1
+            elif order < self.max_order:
+                error_kp1 = _compute_error_ratio(
+                    tuple(dt_cast * gamma_star[order] * iphi_ for iphi_ in implicit_phi_p[order]), tolerance
+                )
+                if max(error_kp1) < max(error_k):
+                    next_order = order + 1
+
+        # Keep step size constant if increasing order. Else use adaptive step size.
+        dt_next = dt if next_order > order else _optimal_step_size(
+            dt, error_k, self.safety, self.ifactor, self.dfactor, order=order + 1
+        )
+
+        prev_f.appendleft(next_f0)
+        prev_t.appendleft(next_t)
+        return _VCABMState(p_next, prev_f, prev_t, next_t + dt_next, implicit_phi, order=next_order)

ShapeID/DiffEqs/adjoint.py

@@ -0,0 +1,133 @@
+import torch
+import torch.nn as nn
+from ShapeID.DiffEqs.odeint import odeint
+from ShapeID.DiffEqs.misc import _flatten, _flatten_convert_none_to_zeros
+
+
+class OdeintAdjointMethod(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, *args):
+        assert len(args) >= 8, 'Internal error: all arguments required.'
+        y0, func, t, dt, flat_params, rtol, atol, method, options = \
+            args[:-8], args[-8], args[-7], args[-6], args[-5], args[-4], args[-3], args[-2], args[-1]
+
+        ctx.func, ctx.rtol, ctx.atol, ctx.method, ctx.options = func, rtol, atol, method, options
+
+        with torch.no_grad():
+            ans = odeint(func, y0, t, dt, rtol=rtol, atol=atol, method=method, options=options)
+            ctx.save_for_backward(t, flat_params, *ans)
+        return ans
+
+    @staticmethod
+    def backward(ctx, *grad_output):
+
+        t, flat_params, *ans = ctx.saved_tensors
+        ans = tuple(ans)
+        func, rtol, atol, method, options = ctx.func, ctx.rtol, ctx.atol, ctx.method, ctx.options
+        n_tensors = len(ans)
+        f_params = tuple(func.parameters())
+
+        # TODO: use a nn.Module and call odeint_adjoint to implement higher order derivatives.
+        def augmented_dynamics(t, y_aug):
+            # Dynamics of the original system augmented with
+            # the adjoint wrt y, and an integrator wrt t and args.
+            y, adj_y = y_aug[:n_tensors], y_aug[n_tensors:2 * n_tensors]  # Ignore adj_time and adj_params.
+
+            with torch.set_grad_enabled(True):
+                t = t.to(y[0].device).detach().requires_grad_(True)
+                y = tuple(y_.detach().requires_grad_(True) for y_ in y)
+                func_eval = func(t, y)
+                vjp_t, *vjp_y_and_params = torch.autograd.grad(
+                    func_eval, (t,) + y + f_params,
+                    tuple(-adj_y_ for adj_y_ in adj_y), allow_unused=True, retain_graph=True
+                )
+            vjp_y = vjp_y_and_params[:n_tensors]
+            vjp_params = vjp_y_and_params[n_tensors:]
+
+            # autograd.grad returns None if no gradient, set to zero.
+            vjp_t = torch.zeros_like(t) if vjp_t is None else vjp_t
+            vjp_y = tuple(torch.zeros_like(y_) if vjp_y_ is None else vjp_y_ for vjp_y_, y_ in zip(vjp_y, y))
+            vjp_params = _flatten_convert_none_to_zeros(vjp_params, f_params)
+
+            if len(f_params) == 0:
+                vjp_params = torch.tensor(0.).to(vjp_y[0])
+            return (*func_eval, *vjp_y, vjp_t, vjp_params)
+
+        T = ans[0].shape[0]
+        with torch.no_grad():
+            adj_y = tuple(grad_output_[-1] for grad_output_ in grad_output)
+            adj_params = torch.zeros_like(flat_params)
+            adj_time = torch.tensor(0.).to(t)
+            time_vjps = []
+            for i in range(T - 1, 0, -1):
+
+                ans_i = tuple(ans_[i] for ans_ in ans)
+                grad_output_i = tuple(grad_output_[i] for grad_output_ in grad_output)
+                func_i = func(t[i], ans_i)
+
+                # Compute the effect of moving the current time measurement point.
+                dLd_cur_t = sum(
+                    torch.dot(func_i_.reshape(-1), grad_output_i_.reshape(-1)).reshape(1)
+                    for func_i_, grad_output_i_ in zip(func_i, grad_output_i)
+                )
+                adj_time = adj_time - dLd_cur_t
+                time_vjps.append(dLd_cur_t)
+
+                # Run the augmented system backwards in time.
+                if adj_params.numel() == 0:
+                    adj_params = torch.tensor(0.).to(adj_y[0])
+                aug_y0 = (*ans_i, *adj_y, adj_time, adj_params)
+                aug_ans = odeint(
+                    augmented_dynamics, aug_y0,
+                    torch.tensor([t[i], t[i - 1]]), rtol=rtol, atol=atol, method=method, options=options
+                )
+
+                # Unpack aug_ans.
+                adj_y = aug_ans[n_tensors:2 * n_tensors]
+                adj_time = aug_ans[2 * n_tensors]
+                adj_params = aug_ans[2 * n_tensors + 1]
+
+                adj_y = tuple(adj_y_[1] if len(adj_y_) > 0 else adj_y_ for adj_y_ in adj_y)
+                if len(adj_time) > 0: adj_time = adj_time[1]
+                if len(adj_params) > 0: adj_params = adj_params[1]
+
+                adj_y = tuple(adj_y_ + grad_output_[i - 1] for adj_y_, grad_output_ in zip(adj_y, grad_output))
+
+                del aug_y0, aug_ans
+
+            time_vjps.append(adj_time)
+            time_vjps = torch.cat(time_vjps[::-1])
+
+            return (*adj_y, None, time_vjps, adj_params, None, None, None, None, None, None) # Add a None (TODO, futher check)
+
+
+def odeint_adjoint(func, y0, t, dt, rtol=1e-6, atol=1e-12, method=None, options=None):
+
+    # We need this in order to access the variables inside this module,
+    # since we have no other way of getting variables along the execution path.
+    if not isinstance(func, nn.Module):
+        raise ValueError('func is required to be an instance of nn.Module.')
+
+    tensor_input = False
+    if torch.is_tensor(y0):
+
+        class TupleFunc(nn.Module):
+
+            def __init__(self, base_func):
+                super(TupleFunc, self).__init__()
+                self.base_func = base_func
+
+            def forward(self, t, y):
+                return (self.base_func(t, y[0]),)
+
+        tensor_input = True
+        y0 = (y0,)
+        func = TupleFunc(func)
+
+    flat_params = _flatten(func.parameters())
+    ys = OdeintAdjointMethod.apply(*y0, func, t, dt, flat_params, rtol, atol, method, options)
+
+    if tensor_input:
+        ys = ys[0]
+    return ys

ShapeID/DiffEqs/dopri5.py

@@ -0,0 +1,172 @@
+import torch
+from .misc import (
+    _scaled_dot_product, _convert_to_tensor, _is_finite, _select_initial_step, _handle_unused_kwargs, _is_iterable,
+    _optimal_step_size, _compute_error_ratio
+)
+from .solvers import AdaptiveStepsizeODESolver, set_BC_2D, set_BC_3D, add_dBC_2D, add_dBC_3D
+from .interp import _interp_fit, _interp_evaluate
+from .rk_common import _RungeKuttaState, _ButcherTableau, _runge_kutta_step
+
+
+_DORMAND_PRINCE_SHAMPINE_TABLEAU = _ButcherTableau(
+    alpha=[1 / 5, 3 / 10, 4 / 5, 8 / 9, 1., 1.],
+    beta=[
+        [1 / 5],
+        [3 / 40, 9 / 40],
+        [44 / 45, -56 / 15, 32 / 9],
+        [19372 / 6561, -25360 / 2187, 64448 / 6561, -212 / 729],
+        [9017 / 3168, -355 / 33, 46732 / 5247, 49 / 176, -5103 / 18656],
+        [35 / 384, 0, 500 / 1113, 125 / 192, -2187 / 6784, 11 / 84],
+    ],
+    c_sol=[35 / 384, 0, 500 / 1113, 125 / 192, -2187 / 6784, 11 / 84, 0],
+    c_error=[
+        35 / 384 - 1951 / 21600,
+        0,
+        500 / 1113 - 22642 / 50085,
+        125 / 192 - 451 / 720,
+        -2187 / 6784 - -12231 / 42400,
+        11 / 84 - 649 / 6300,
+        -1. / 60.,
+    ],
+)
+
+DPS_C_MID = [
+    6025192743 / 30085553152 / 2, 0, 51252292925 / 65400821598 / 2, -2691868925 / 45128329728 / 2,
+    187940372067 / 1594534317056 / 2, -1776094331 / 19743644256 / 2, 11237099 / 235043384 / 2
+]
+
+
+def _interp_fit_dopri5(y0, y1, k, dt, tableau=_DORMAND_PRINCE_SHAMPINE_TABLEAU):
+    """Fit an interpolating polynomial to the results of a Runge-Kutta step."""
+    dt = dt.type_as(y0[0])
+    y_mid = tuple(y0_ + _scaled_dot_product(dt, DPS_C_MID, k_) for y0_, k_ in zip(y0, k))
+    f0 = tuple(k_[0] for k_ in k)
+    f1 = tuple(k_[-1] for k_ in k)
+    return _interp_fit(y0, y1, y_mid, f0, f1, dt)
+
+
+def _abs_square(x):
+    return torch.mul(x, x)
+
+
+def _ta_append(list_of_tensors, value):
+    """Append a value to the end of a list of PyTorch tensors."""
+    list_of_tensors.append(value)
+    return list_of_tensors
+
+
+class Dopri5Solver(AdaptiveStepsizeODESolver):
+
+    def __init__(
+        self, func, y0, rtol, atol, dt, first_step=None, safety=0.9, ifactor=10.0, dfactor=0.2, max_num_steps=2**31 - 1,
+        options = None
+        #**unused_kwargs
+    ):
+        #_handle_unused_kwargs(self, unused_kwargs)
+        #del unused_kwargs
+
+        self.func = func
+        self.y0 = y0
+
+        self.dt = dt #options.dt
+        '''if 'dirichlet' in options.BC or 'cauchy' in options.BC and options.contours is not None:
+            self.contours = options.contours  # (n_batch, nT, 4 / 6, BC_size, sub_spatial_shape)
+            self.BC_size = self.contours.size(3)
+            self.set_BC = set_BC_2D if self.contours.size(2) == 4 else set_BC_3D
+        else:
+            self.contours = None
+        if 'source' in options.BC and options.dcontours is not None:
+            self.dcontours = options.dcontours  # (n_batch, nT, 4 / 6, BC_size, sub_spatial_shape)
+            self.BC_size = self.dcontours.size(3)
+            self.add_dBC = add_dBC_2D if self.dcontours.size(2) == 4 else add_dBC_3D
+        else:
+            self.dcontours = None'''
+        
+        #self.adjoint = options.adjoint
+
+        self.rtol = rtol if _is_iterable(rtol) else [rtol] * len(y0)
+        self.atol = atol if _is_iterable(atol) else [atol] * len(y0)
+        self.first_step = first_step
+        self.safety = _convert_to_tensor(safety, dtype=torch.float64, device=y0[0].device)
+        self.ifactor = _convert_to_tensor(ifactor, dtype=torch.float64, device=y0[0].device)
+        self.dfactor = _convert_to_tensor(dfactor, dtype=torch.float64, device=y0[0].device)
+        self.max_num_steps = _convert_to_tensor(max_num_steps, dtype=torch.int32, device=y0[0].device)
+        #self.n_step_record=[]
+
+    def before_integrate(self, t):
+        f0 = self.func(t[0].type_as(self.y0[0]), self.y0)
+        #print("first_step is {}".format(self.first_step))
+        if self.first_step is None:
+            first_step = _select_initial_step(self.func, t[0], self.y0, 4, self.rtol[0], self.atol[0], f0=f0).to(t)
+        else:
+            first_step = _convert_to_tensor(0.01, dtype=t.dtype, device=t.device)
+        # if first_step>0.2:
+        #     print("warning the first step of dopri5 {} is too big, set to 0.2".format(first_step))
+        #     first_step = _convert_to_tensor(0.2, dtype=torch.float64, device=self.y0[0].device)
+
+        self.rk_state = _RungeKuttaState(self.y0, f0, t[0], t[0], first_step, interp_coeff=[self.y0] * 5)
+
+    def advance(self, next_t):
+        """Interpolate through the next time point, integrating as necessary."""
+        n_steps = 0
+        while next_t > self.rk_state.t1:
+            assert n_steps < self.max_num_steps, 'max_num_steps exceeded ({}>={})'.format(n_steps, self.max_num_steps)
+            self.rk_state = self._adaptive_dopri5_step(self.rk_state)
+            n_steps += 1
+        # if len(self.n_step_record)==100:
+        #     print("this dopri5 step info will print every 100 calls, the current average step is {}".format(sum(self.n_step_record)/100))
+        #     self.n_step_record=[]
+        # else:
+        #     self.n_step_record.append(n_steps)
+
+        return _interp_evaluate(self.rk_state.interp_coeff, self.rk_state.t0, self.rk_state.t1, next_t)
+
+    def _adaptive_dopri5_step(self, rk_state):
+        """Take an adaptive Runge-Kutta step to integrate the DiffEqs."""
+        y0, f0, _, t0, dt, interp_coeff = rk_state
+        ########################################################
+        #                      Assertions                      #
+        ########################################################
+        assert t0 + dt > t0, 'underflow in dt {}'.format(dt.item())
+        # for y0_ in y0:
+        #     #assert _is_finite(torch.abs(y0_)), 'non-finite values in state `y`: {}'.format(y0_)
+        #     is_finite= _is_finite(torch.abs(y0_))
+        #     if not is_finite:
+        #         print(" non-finite elements exist, try to fix")
+        #         y0_[y0_ != y0_] = 0.
+        #         y0_[y0_ == float("Inf")] = 0.
+
+        y1, f1, y1_error, k = _runge_kutta_step(self.func, y0, f0, t0, dt, tableau=_DORMAND_PRINCE_SHAMPINE_TABLEAU)
+
+        ########################################################
+        #                     Error Ratio                      #
+        ########################################################
+        mean_sq_error_ratio = _compute_error_ratio(y1_error, atol=self.atol, rtol=self.rtol, y0=y0, y1=y1)
+        accept_step = (torch.tensor(mean_sq_error_ratio) <= 1).all()
+
+        ########################################################
+        #                   Update RK State                    #
+        ########################################################
+        dt_next = _optimal_step_size(
+            dt, mean_sq_error_ratio, safety=self.safety, ifactor=self.ifactor, dfactor=self.dfactor, order=5)
+        tol_min_dt = 0.2 * self.dt if 0.1 * self.dt >= 0.01 else 0.01
+        #print('tol min', tol_min_dt)
+        if not (dt_next< tol_min_dt or dt_next>0.1): #(dt_next<0.01 or dt_next>0.1):  #(dt_next<0.02): #not (dt_next<0.02 or dt_next>0.1):
+            y_next = y1 if accept_step else y0
+            f_next = f1 if accept_step else f0
+            t_next = t0 + dt if accept_step else t0
+            interp_coeff = _interp_fit_dopri5(y0, y_next, k, dt) if accept_step else interp_coeff
+        else:
+            if dt_next< tol_min_dt: #dt_next<0.01: # 0.01
+                #print("Dopri5 step %.3f too small, set to %.3f" % (dt_next, 0.2 * self.dt))
+                dt_next = _convert_to_tensor(tol_min_dt, dtype=torch.float64, device=y0[0].device)
+            if dt_next>0.1:
+                #print("Dopri5 step %.8f is too big, set to 0.1" % (dt_next))
+                dt_next = _convert_to_tensor(0.1, dtype=torch.float64, device=y0[0].device)
+            y_next = y1
+            f_next = f1
+            t_next = t0 + dt
+            interp_coeff = _interp_fit_dopri5(y0, y1, k, dt)
+        rk_state = _RungeKuttaState(y_next, f_next, t0, t_next, dt_next, interp_coeff)
+        #print('dt_next', dt_next)
+        return rk_state

ShapeID/DiffEqs/fixed_adams.py

@@ -0,0 +1,211 @@
+import sys
+import collections
+from ShapeID.DiffEqs.solvers import FixedGridODESolver
+from ShapeID.DiffEqs.misc import _scaled_dot_product, _has_converged
+import ShapeID.DiffEqs.rk_common
+
+_BASHFORTH_COEFFICIENTS = [
+    [],  # order 0
+    [11],
+    [3, -1],
+    [23, -16, 5],
+    [55, -59, 37, -9],
+    [1901, -2774, 2616, -1274, 251],
+    [4277, -7923, 9982, -7298, 2877, -475],
+    [198721, -447288, 705549, -688256, 407139, -134472, 19087],
+    [434241, -1152169, 2183877, -2664477, 2102243, -1041723, 295767, -36799],
+    [14097247, -43125206, 95476786, -139855262, 137968480, -91172642, 38833486, -9664106, 1070017],
+    [30277247, -104995189, 265932680, -454661776, 538363838, -444772162, 252618224, -94307320, 20884811, -2082753],
+    [
+        2132509567, -8271795124, 23591063805, -46113029016, 63716378958, -63176201472, 44857168434, -22329634920,
+        7417904451, -1479574348, 134211265
+    ],
+    [
+        4527766399, -19433810163, 61633227185, -135579356757, 214139355366, -247741639374, 211103573298, -131365867290,
+        58189107627, -17410248271, 3158642445, -262747265
+    ],
+    [
+        13064406523627, -61497552797274, 214696591002612, -524924579905150, 932884546055895, -1233589244941764,
+        1226443086129408, -915883387152444, 507140369728425, -202322913738370, 55060974662412, -9160551085734,
+        703604254357
+    ],
+    [
+        27511554976875, -140970750679621, 537247052515662, -1445313351681906, 2854429571790805, -4246767353305755,
+        4825671323488452, -4204551925534524, 2793869602879077, -1393306307155755, 505586141196430, -126174972681906,
+        19382853593787, -1382741929621
+    ],
+    [
+        173233498598849, -960122866404112, 3966421670215481, -11643637530577472, 25298910337081429, -41825269932507728,
+        53471026659940509, -53246738660646912, 41280216336284259, -24704503655607728, 11205849753515179,
+        -3728807256577472, 859236476684231, -122594813904112, 8164168737599
+    ],
+    [
+        362555126427073, -2161567671248849, 9622096909515337, -30607373860520569, 72558117072259733,
+        -131963191940828581, 187463140112902893, -210020588912321949, 186087544263596643, -129930094104237331,
+        70724351582843483, -29417910911251819, 9038571752734087, -1934443196892599, 257650275915823, -16088129229375
+    ],
+    [
+        192996103681340479, -1231887339593444974, 5878428128276811750, -20141834622844109630, 51733880057282977010,
+        -102651404730855807942, 160414858999474733422, -199694296833704562550, 199061418623907202560,
+        -158848144481581407370, 100878076849144434322, -50353311405771659322, 19338911944324897550,
+        -5518639984393844930, 1102560345141059610, -137692773163513234, 8092989203533249
+    ],
+    [
+        401972381695456831, -2735437642844079789, 13930159965811142228, -51150187791975812900, 141500575026572531760,
+        -304188128232928718008, 518600355541383671092, -710171024091234303204, 786600875277595877750,
+        -706174326992944287370, 512538584122114046748, -298477260353977522892, 137563142659866897224,
+        -49070094880794267600, 13071639236569712860, -2448689255584545196, 287848942064256339, -15980174332775873
+    ],
+    [
+        333374427829017307697, -2409687649238345289684, 13044139139831833251471, -51099831122607588046344,
+        151474888613495715415020, -350702929608291455167896, 647758157491921902292692, -967713746544629658690408,
+        1179078743786280451953222, -1176161829956768365219840, 960377035444205950813626, -639182123082298748001432,
+        343690461612471516746028, -147118738993288163742312, 48988597853073465932820, -12236035290567356418552,
+        2157574942881818312049, -239560589366324764716, 12600467236042756559
+    ],
+    [
+        691668239157222107697, -5292843584961252933125, 30349492858024727686755, -126346544855927856134295,
+        399537307669842150996468, -991168450545135070835076, 1971629028083798845750380, -3191065388846318679544380,
+        4241614331208149947151790, -4654326468801478894406214, 4222756879776354065593786, -3161821089800186539248210,
+        1943018818982002395655620, -970350191086531368649620, 387739787034699092364924, -121059601023985433003532,
+        28462032496476316665705, -4740335757093710713245, 498669220956647866875, -24919383499187492303
+    ],
+]
+
+_MOULTON_COEFFICIENTS = [
+    [],  # order 0
+    [1],
+    [1, 1],
+    [5, 8, -1],
+    [9, 19, -5, 1],
+    [251, 646, -264, 106, -19],
+    [475, 1427, -798, 482, -173, 27],
+    [19087, 65112, -46461, 37504, -20211, 6312, -863],
+    [36799, 139849, -121797, 123133, -88547, 41499, -11351, 1375],
+    [1070017, 4467094, -4604594, 5595358, -5033120, 3146338, -1291214, 312874, -33953],
+    [2082753, 9449717, -11271304, 16002320, -17283646, 13510082, -7394032, 2687864, -583435, 57281],
+    [
+        134211265, 656185652, -890175549, 1446205080, -1823311566, 1710774528, -1170597042, 567450984, -184776195,
+        36284876, -3250433
+    ],
+    [
+        262747265, 1374799219, -2092490673, 3828828885, -5519460582, 6043521486, -4963166514, 3007739418, -1305971115,
+        384709327, -68928781, 5675265
+    ],
+    [
+        703604254357, 3917551216986, -6616420957428, 13465774256510, -21847538039895, 27345870698436, -26204344465152,
+        19058185652796, -10344711794985, 4063327863170, -1092096992268, 179842822566, -13695779093
+    ],
+    [
+        1382741929621, 8153167962181, -15141235084110, 33928990133618, -61188680131285, 86180228689563, -94393338653892,
+        80101021029180, -52177910882661, 25620259777835, -9181635605134, 2268078814386, -345457086395, 24466579093
+    ],
+    [
+        8164168737599, 50770967534864, -102885148956217, 251724894607936, -499547203754837, 781911618071632,
+        -963605400824733, 934600833490944, -710312834197347, 418551804601264, -187504936597931, 61759426692544,
+        -14110480969927, 1998759236336, -132282840127
+    ],
+    [
+        16088129229375, 105145058757073, -230992163723849, 612744541065337, -1326978663058069, 2285168598349733,
+        -3129453071993581, 3414941728852893, -2966365730265699, 2039345879546643, -1096355235402331, 451403108933483,
+        -137515713789319, 29219384284087, -3867689367599, 240208245823
+    ],
+    [
+        8092989203533249, 55415287221275246, -131240807912923110, 375195469874202430, -880520318434977010,
+        1654462865819232198, -2492570347928318318, 3022404969160106870, -2953729295811279360, 2320851086013919370,
+        -1455690451266780818, 719242466216944698, -273894214307914510, 77597639915764930, -15407325991235610,
+        1913813460537746, -111956703448001
+    ],
+    [
+        15980174332775873, 114329243705491117, -290470969929371220, 890337710266029860, -2250854333681641520,
+        4582441343348851896, -7532171919277411636, 10047287575124288740, -10910555637627652470, 9644799218032932490,
+        -6913858539337636636, 3985516155854664396, -1821304040326216520, 645008976643217360, -170761422500096220,
+        31816981024600492, -3722582669836627, 205804074290625
+    ],
+    [
+        12600467236042756559, 93965550344204933076, -255007751875033918095, 834286388106402145800,
+        -2260420115705863623660, 4956655592790542146968, -8827052559979384209108, 12845814402199484797800,
+        -15345231910046032448070, 15072781455122686545920, -12155867625610599812538, 8008520809622324571288,
+        -4269779992576330506540, 1814584564159445787240, -600505972582990474260, 149186846171741510136,
+        -26182538841925312881, 2895045518506940460, -151711881512390095
+    ],
+    [
+        24919383499187492303, 193280569173472261637, -558160720115629395555, 1941395668950986461335,
+        -5612131802364455926260, 13187185898439270330756, -25293146116627869170796, 39878419226784442421820,
+        -51970649453670274135470, 56154678684618739939910, -50320851025594566473146, 37297227252822858381906,
+        -22726350407538133839300, 11268210124987992327060, -4474886658024166985340, 1389665263296211699212,
+        -325187970422032795497, 53935307402575440285, -5652892248087175675, 281550972898020815
+    ],
+]
+
+_DIVISOR = [
+    None, 11, 2, 12, 24, 720, 1440, 60480, 120960, 3628800, 7257600, 479001600, 958003200, 2615348736000, 5230697472000,
+    31384184832000, 62768369664000, 32011868528640000, 64023737057280000, 51090942171709440000, 102181884343418880000
+]
+
+_MIN_ORDER = 4
+_MAX_ORDER = 12
+_MAX_ITERS = 4
+
+
+class AdamsBashforthMoulton(FixedGridODESolver):
+
+    def __init__(
+        self, func, y0, rtol=1e-3, atol=1e-4, implicit=True, max_iters=_MAX_ITERS, max_order=_MAX_ORDER, **kwargs
+    ):
+        super(AdamsBashforthMoulton, self).__init__(func, y0, **kwargs)
+
+        self.rtol = rtol
+        self.atol = atol
+        self.implicit = implicit
+        self.max_iters = max_iters
+        self.max_order = int(min(max_order, _MAX_ORDER))
+        self.prev_f = collections.deque(maxlen=self.max_order - 1)
+        self.prev_t = None
+
+    def _update_history(self, t, f):
+        if self.prev_t is None or self.prev_t != t:
+            self.prev_f.appendleft(f)
+            self.prev_t = t
+
+    def step_func(self, func, t, dt, y):
+        self._update_history(t, func(t, y))
+        order = min(len(self.prev_f), self.max_order - 1)
+        if order < _MIN_ORDER - 1:
+            # Compute using RK4.
+            dy = rk_common.rk4_alt_step_func(func, t, dt, y, k1=self.prev_f[0])
+            return dy
+        else:
+            # Adams-Bashforth predictor.
+            bashforth_coeffs = _BASHFORTH_COEFFICIENTS[order]
+            ab_div = _DIVISOR[order]
+            dy = tuple(dt * _scaled_dot_product(1 / ab_div, bashforth_coeffs, f_) for f_ in zip(*self.prev_f))
+
+            # Adams-Moulton corrector.
+            if self.implicit:
+                moulton_coeffs = _MOULTON_COEFFICIENTS[order + 1]
+                am_div = _DIVISOR[order + 1]
+                delta = tuple(dt * _scaled_dot_product(1 / am_div, moulton_coeffs[1:], f_) for f_ in zip(*self.prev_f))
+                converged = False
+                for _ in range(self.max_iters):
+                    dy_old = dy
+                    f = func(t + dt, tuple(y_ + dy_ for y_, dy_ in zip(y, dy)))
+                    dy = tuple(dt * (moulton_coeffs[0] / am_div) * f_ + delta_ for f_, delta_ in zip(f, delta))
+                    converged = _has_converged(dy_old, dy, self.rtol, self.atol)
+                    if converged:
+                        break
+                if not converged:
+                    print('Warning: Functional iteration did not converge. Solution may be incorrect.', file=sys.stderr)
+                    self.prev_f.pop()
+                self._update_history(t, f)
+            return dy
+
+    @property
+    def order(self):
+        return 4
+
+
+class AdamsBashforth(AdamsBashforthMoulton):
+
+    def __init__(self, func, y0, **kwargs):
+        super(AdamsBashforth, self).__init__(func, y0, implicit=False, **kwargs)

ShapeID/DiffEqs/fixed_grid.py

@@ -0,0 +1,33 @@
+from ShapeID.DiffEqs.solvers import FixedGridODESolver
+import ShapeID.DiffEqs.rk_common as rk_common
+
+
+class Euler(FixedGridODESolver):
+
+    def step_func(self, func, t, dt, y):
+        return tuple(dt * f_ for f_ in func(t, y))
+
+    @property
+    def order(self):
+        return 1
+
+
+class Midpoint(FixedGridODESolver):
+
+    def step_func(self, func, t, dt, y):
+        y_mid = tuple(y_ + f_ * dt / 2 for y_, f_ in zip(y, func(t, y)))
+        return tuple(dt * f_ for f_ in func(t + dt / 2, y_mid))
+
+    @property
+    def order(self):
+        return 2
+
+
+class RK4(FixedGridODESolver):
+
+    def step_func(self, func, t, dt, y):
+        return rk_common.rk4_alt_step_func(func, t, dt, y)
+
+    @property
+    def order(self):
+        return 4

ShapeID/DiffEqs/interp.py

@@ -0,0 +1,65 @@
+import torch
+from ShapeID.DiffEqs.misc import _convert_to_tensor, _dot_product
+
+
+def _interp_fit(y0, y1, y_mid, f0, f1, dt):
+    """Fit coefficients for 4th order polynomial interpolation.
+
+    Args:
+        y0: function value at the start of the interval.
+        y1: function value at the end of the interval.
+        y_mid: function value at the mid-point of the interval.
+        f0: derivative value at the start of the interval.
+        f1: derivative value at the end of the interval.
+        dt: width of the interval.
+
+    Returns:
+        List of coefficients `[a, b, c, d, e]` for interpolating with the polynomial
+        `p = a * x ** 4 + b * x ** 3 + c * x ** 2 + d * x + e` for values of `x`
+        between 0 (start of interval) and 1 (end of interval).
+    """
+    a = tuple(
+        _dot_product([-2 * dt, 2 * dt, -8, -8, 16], [f0_, f1_, y0_, y1_, y_mid_])
+        for f0_, f1_, y0_, y1_, y_mid_ in zip(f0, f1, y0, y1, y_mid)
+    )
+    b = tuple(
+        _dot_product([5 * dt, -3 * dt, 18, 14, -32], [f0_, f1_, y0_, y1_, y_mid_])
+        for f0_, f1_, y0_, y1_, y_mid_ in zip(f0, f1, y0, y1, y_mid)
+    )
+    c = tuple(
+        _dot_product([-4 * dt, dt, -11, -5, 16], [f0_, f1_, y0_, y1_, y_mid_])
+        for f0_, f1_, y0_, y1_, y_mid_ in zip(f0, f1, y0, y1, y_mid)
+    )
+    d = tuple(dt * f0_ for f0_ in f0)
+    e = y0
+    return [a, b, c, d, e]
+
+
+def _interp_evaluate(coefficients, t0, t1, t):
+    """Evaluate polynomial interpolation at the given time point.
+
+    Args:
+        coefficients: list of Tensor coefficients as created by `interp_fit`.
+        t0: scalar float64 Tensor giving the start of the interval.
+        t1: scalar float64 Tensor giving the end of the interval.
+        t: scalar float64 Tensor giving the desired interpolation point.
+
+    Returns:
+        Polynomial interpolation of the coefficients at time `t`.
+    """
+
+    dtype = coefficients[0][0].dtype
+    device = coefficients[0][0].device
+
+    t0 = _convert_to_tensor(t0, dtype=dtype, device=device)
+    t1 = _convert_to_tensor(t1, dtype=dtype, device=device)
+    t = _convert_to_tensor(t, dtype=dtype, device=device)
+
+    assert (t0 <= t) & (t <= t1), 'invalid interpolation, fails `t0 <= t <= t1`: {}, {}, {}'.format(t0, t, t1)
+    x = ((t - t0) / (t1 - t0)).type(dtype).to(device)
+
+    xs = [torch.tensor(1).type(dtype).to(device), x]
+    for _ in range(2, len(coefficients)):
+        xs.append(xs[-1] * x)
+
+    return tuple(_dot_product(coefficients_, reversed(xs)) for coefficients_ in zip(*coefficients))

ShapeID/DiffEqs/misc.py

@@ -0,0 +1,195 @@
+import warnings
+import torch
+
+
+def _flatten(sequence):
+    flat = [p.contiguous().view(-1) for p in sequence]
+    return torch.cat(flat) if len(flat) > 0 else torch.tensor([])
+
+
+def _flatten_convert_none_to_zeros(sequence, like_sequence):
+    flat = [
+        p.contiguous().view(-1) if p is not None else torch.zeros_like(q).view(-1)
+        for p, q in zip(sequence, like_sequence)
+    ]
+    return torch.cat(flat) if len(flat) > 0 else torch.tensor([])
+
+
+def _possibly_nonzero(x):
+    return isinstance(x, torch.Tensor) or x != 0
+
+
+def _scaled_dot_product(scale, xs, ys):
+    """Calculate a scaled, vector inner product between lists of Tensors."""
+    # Using _possibly_nonzero lets us avoid wasted computation.
+    return sum([(scale * x) * y for x, y in zip(xs, ys) if _possibly_nonzero(x) or _possibly_nonzero(y)])
+
+
+def _dot_product(xs, ys):
+    """Calculate the vector inner product between two lists of Tensors."""
+    return sum([x * y for x, y in zip(xs, ys)])
+
+
+def _has_converged(y0, y1, rtol, atol):
+    """Checks that each element is within the error tolerance."""
+    error_tol = tuple(atol + rtol * torch.max(torch.abs(y0_), torch.abs(y1_)) for y0_, y1_ in zip(y0, y1))
+    error = tuple(torch.abs(y0_ - y1_) for y0_, y1_ in zip(y0, y1))
+    return all((error_ < error_tol_).all() for error_, error_tol_ in zip(error, error_tol))
+
+
+def _convert_to_tensor(a, dtype=None, device=None):
+    if not isinstance(a, torch.Tensor):
+        a = torch.tensor(a)
+    if dtype is not None:
+        a = a.type(dtype)
+    if device is not None:
+        a = a.to(device)
+    return a
+
+
+def _is_finite(tensor):
+    _check = (tensor == float('inf')) + (tensor == float('-inf')) + torch.isnan(tensor)
+    return not _check.any()
+
+
+def _decreasing(t):
+    return (t[1:] < t[:-1]).all()
+
+
+def _assert_increasing(t):  
+    assert (t[1:] > t[:-1]).all(), 't must be strictly increasing or decrasing'
+
+
+def _is_iterable(inputs):
+    try:
+        iter(inputs)
+        return True
+    except TypeError:
+        return False
+
+
+def _norm(x):
+    """Compute RMS norm."""
+    if torch.is_tensor(x):
+        return x.norm() / (x.numel()**0.5)
+    else:
+        return torch.sqrt(sum(x_.norm()**2 for x_ in x) / sum(x_.numel() for x_ in x))
+
+
+def _handle_unused_kwargs(solver, unused_kwargs):
+    if len(unused_kwargs) > 0:
+        warnings.warn('{}: Unexpected arguments {}'.format(solver.__class__.__name__, unused_kwargs))
+
+
+def _select_initial_step(fun, t0, y0, order, rtol, atol, f0=None):
+    """Empirically select a good initial step.
+
+    The algorithm is described in [1]_.
+
+    Parameters
+    ----------
+    fun : callable
+        Right-hand side of the system.
+    t0 : float
+        Initial value of the independent variable.
+    y0 : ndarray, shape (n,)
+        Initial value of the dependent variable.
+    direction : float
+        Integration direction.
+    order : float
+        Method order.
+    rtol : float
+        Desired relative tolerance.
+    atol : float
+        Desired absolute tolerance.
+
+    Returns
+    -------
+    h_abs : float
+        Absolute value of the suggested initial step.
+
+    References
+    ----------
+    .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential
+           Equations I: Nonstiff Problems", Sec. II.4.
+    """
+    t0 = t0.to(y0[0])
+    if f0 is None:
+        f0 = fun(t0, y0)
+
+    rtol = rtol if _is_iterable(rtol) else [rtol] * len(y0)
+    atol = atol if _is_iterable(atol) else [atol] * len(y0)
+
+    scale = tuple(atol_ + torch.abs(y0_) * rtol_ for y0_, atol_, rtol_ in zip(y0, atol, rtol))
+
+    d0 = tuple(_norm(y0_ / scale_) for y0_, scale_ in zip(y0, scale))
+    d1 = tuple(_norm(f0_ / scale_) for f0_, scale_ in zip(f0, scale))
+
+    if max(d0).item() < 1e-5 or max(d1).item() < 1e-5:
+        h0 = torch.tensor(1e-6).to(t0)
+    else:
+        h0 = 0.01 * max(d0_ / d1_ for d0_, d1_ in zip(d0, d1))
+
+    y1 = tuple(y0_ + h0 * f0_ for y0_, f0_ in zip(y0, f0))
+    f1 = fun(t0 + h0, y1)
+
+    d2 = tuple(_norm((f1_ - f0_) / scale_) / h0 for f1_, f0_, scale_ in zip(f1, f0, scale))
+
+    if max(d1).item() <= 1e-15 and max(d2).item() <= 1e-15:
+        h1 = torch.max(torch.tensor(1e-6).to(h0), h0 * 1e-3)
+    else:
+        h1 = (0.01 / max(d1 + d2))**(1. / float(order + 1))
+
+    return torch.min(100 * h0, h1)
+
+
+def _compute_error_ratio(error_estimate, error_tol=None, rtol=None, atol=None, y0=None, y1=None):
+    if error_tol is None:
+        assert rtol is not None and atol is not None and y0 is not None and y1 is not None
+        rtol if _is_iterable(rtol) else [rtol] * len(y0)
+        atol if _is_iterable(atol) else [atol] * len(y0)
+        error_tol = tuple(
+            atol_ + rtol_ * torch.max(torch.abs(y0_), torch.abs(y1_))
+            for atol_, rtol_, y0_, y1_ in zip(atol, rtol, y0, y1)
+        )
+    error_ratio = tuple(error_estimate_ / error_tol_ for error_estimate_, error_tol_ in zip(error_estimate, error_tol))
+    mean_sq_error_ratio = tuple(torch.mean(error_ratio_ * error_ratio_) for error_ratio_ in error_ratio)
+    return mean_sq_error_ratio
+
+
+def _optimal_step_size(last_step, mean_error_ratio, safety=0.9, ifactor=10.0, dfactor=0.2, order=5):
+    """Calculate the optimal size for the next step."""
+    mean_error_ratio = max(mean_error_ratio)  # Compute step size based on highest ratio.
+    if mean_error_ratio == 0:
+        return last_step * ifactor
+    if mean_error_ratio < 1:
+        dfactor = _convert_to_tensor(1, dtype=torch.float64, device=mean_error_ratio.device)
+    error_ratio = torch.sqrt(mean_error_ratio).to(last_step)
+    exponent = torch.tensor(1 / order).to(last_step)
+    factor = torch.max(1 / ifactor, torch.min(error_ratio**exponent / safety, 1 / dfactor))
+    return last_step / factor
+
+
+def _check_inputs(func, y0, t):
+    tensor_input = False
+    if torch.is_tensor(y0):
+        tensor_input = True
+        y0 = (y0,)
+        _base_nontuple_func_ = func
+        func = lambda t, y: (_base_nontuple_func_(t, y[0]),)
+    assert isinstance(y0, tuple), 'y0 must be either a torch.Tensor or a tuple'
+    for y0_ in y0:
+        assert torch.is_tensor(y0_), 'each element must be a torch.Tensor but received {}'.format(type(y0_))
+
+    if _decreasing(t):
+        t = -t
+        _base_reverse_func = func
+        func = lambda t, y: tuple(-f_ for f_ in _base_reverse_func(-t, y))
+
+    for y0_ in y0:
+        if not torch.is_floating_point(y0_):
+            raise TypeError('`y0` must be a floating point Tensor but is a {}'.format(y0_.type()))
+    if not torch.is_floating_point(t):
+        raise TypeError('`t` must be a floating point Tensor but is a {}'.format(t.type()))
+
+    return tensor_input, func, y0, t

ShapeID/DiffEqs/odeint.py

@@ -0,0 +1,75 @@
+from ShapeID.DiffEqs.tsit5 import Tsit5Solver
+from ShapeID.DiffEqs.dopri5 import Dopri5Solver
+from ShapeID.DiffEqs.fixed_grid import Euler, Midpoint, RK4
+from ShapeID.DiffEqs.fixed_adams import AdamsBashforth, AdamsBashforthMoulton
+from ShapeID.DiffEqs.adams import VariableCoefficientAdamsBashforth
+from ShapeID.DiffEqs.misc import _check_inputs
+
+SOLVERS = {
+    'explicit_adams': AdamsBashforth,
+    'fixed_adams': AdamsBashforthMoulton,
+    'adams': VariableCoefficientAdamsBashforth,
+    'tsit5': Tsit5Solver,
+    'dopri5': Dopri5Solver,
+    'euler': Euler,
+    'midpoint': Midpoint,
+    'rk4': RK4,
+}
+
+
+def odeint(func, y0, t, dt, step_size = None, rtol = 1e-7, atol = 1e-9, method = None, options = None):
+    """Integrate a system of ordinary differential equations.
+
+    Solves the initial value problem for a non-stiff system of first order ODEs:
+        ```
+        dy/dt = func(t, y), y(t[0]) = y0
+        ```
+    where y is a Tensor of any shape.
+
+    Output dtypes and numerical precision are based on the dtypes of the inputs `y0`.
+
+    Args:
+        func: Function that maps a Tensor holding the state `y` and a scalar Tensor
+            `t` into a Tensor of state derivatives with respect to time.
+        y0: N-D Tensor giving starting value of `y` at time point `t[0]`. May
+            have any floating point or complex dtype.
+        t: 1-D Tensor holding a sequence of time points for which to solve for
+            `y`. The initial time point should be the first element of this sequence,
+            and each time must be larger than the previous time. May have any floating
+            point dtype. Converted to a Tensor with float64 dtype.
+        rtol: optional float64 Tensor specifying an upper bound on relative error,
+            per element of `y`.
+        atol: optional float64 Tensor specifying an upper bound on absolute error,
+            per element of `y`.
+        method: optional string indicating the integration method to use.
+        options: optional dict of configuring options for the indicated integration
+            method. Can only be provided if a `method` is explicitly set.
+        name: Optional name for this operation.
+
+    Returns:
+        y: Tensor, where the first dimension corresponds to different
+            time points. Contains the solved value of y for each desired time point in
+            `t`, with the initial value `y0` being the first element along the first
+            dimension.
+
+    Raises:
+        ValueError: if an invalid `method` is provided.
+        TypeError: if `options` is supplied without `method`, or if `t` or `y0` has
+            an invalid dtype.
+    """
+
+    tensor_input, func, y0, t = _check_inputs(func, y0, t)
+
+    if options and method is None:
+        raise ValueError('cannot supply `options` without specifying `method`')
+
+    if method is None:
+        method = 'dopri5'
+
+    #solver = SOLVERS[method](func, y0, rtol = rtol, atol = atol, **options)
+    solver = SOLVERS[method](func, y0, rtol = rtol, atol = atol, dt = dt, options = options)
+    solution = solver.integrate(t)
+
+    if tensor_input:
+        solution = solution[0]
+    return solution

ShapeID/DiffEqs/pde.py

@@ -0,0 +1,643 @@
+# ported from https://github.com/pvigier/perlin-numpy
+
+import math
+
+import numpy as np 
+
+import torch
+import torch.nn as nn
+
+
+
+
+def gradient_f(X, batched = False, delta_lst = [1., 1., 1.]):
+    '''
+    Compute gradient of a torch tensor "X" in each direction
+    Upper-boundaries: Backward Difference
+    Non-boundaries & Upper-boundaries: Forward Difference
+    if X is batched: (n_batch, ...);
+    else: (...)
+    '''
+    device = X.device
+    dim = len(X.size()) - 1 if batched else len(X.size())
+    #print(batched)
+    #print(dim)
+    if dim == 1:
+        #print('dim = 1')
+        dX = torch.zeros(X.size(), dtype = torch.float, device = device)
+        X = X.permute(1, 0) if batched else X
+        dX = dX.permute(1, 0) if batched else dX
+        dX[-1] = X[-1] - X[-2] # Backward Difference
+        dX[:-1] = X[1:] - X[:-1] # Forward Difference
+
+        dX = dX.permute(1, 0) if batched else dX
+        dX /= delta_lst[0]
+    elif dim == 2:
+        #print('dim = 2')
+        dX = torch.zeros(X.size() + tuple([2]), dtype = torch.float, device = device)
+        X = X.permute(1, 2, 0) if batched else X
+        dX = dX.permute(1, 2, 3, 0) if batched else dX # put batch to last dim
+        dX[-1, :, 0] = X[-1, :] - X[-2, :] # Backward Difference
+        dX[:-1, :, 0] = X[1:] - X[:-1] # Forward Difference
+
+        dX[:, -1, 1] = X[:, -1] - X[:, -2] # Backward Difference
+        dX[:, :-1, 1] = X[:, 1:] - X[:, :-1] # Forward Difference
+
+        dX = dX.permute(3, 0, 1, 2) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+    elif dim == 3:
+        #print('dim = 3')
+        dX = torch.zeros(X.size() + tuple([3]), dtype = torch.float, device = device)
+        X = X.permute(1, 2, 3, 0) if batched else X
+        dX = dX.permute(1, 2, 3, 4, 0) if batched else dX
+        dX[-1, :, :, 0] = X[-1, :, :] - X[-2, :, :] # Backward Difference
+        dX[:-1, :, :, 0] = X[1:] - X[:-1] # Forward Difference
+
+        dX[:, -1, :, 1] = X[:, -1] - X[:, -2] # Backward Difference
+        dX[:, :-1, :, 1] = X[:, 1:] - X[:, :-1] # Forward Difference
+
+        dX[:, :, -1, 2] = X[:, :, -1] - X[:, :, -2] # Backward Difference
+        dX[:, :, :-1, 2] = X[:, :, 1:] - X[:, :, :-1] # Forward Difference
+
+        dX = dX.permute(4, 0, 1, 2, 3) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+        dX[..., 2] /= delta_lst[2]
+    return dX
+
+
+def gradient_b(X, batched = False, delta_lst = [1., 1., 1.]):
+    '''
+    Compute gradient of a torch tensor "X" in each direction
+    Non-boundaries & Upper-boundaries: Backward Difference
+    Lower-boundaries: Forward Difference
+    if X is batched: (n_batch, ...);
+    else: (...)
+    '''
+    device = X.device
+    dim = len(X.size()) - 1 if batched else len(X.size())
+    #print(batched)
+    #print(dim)
+    if dim == 1:
+        #print('dim = 1')
+        dX = torch.zeros(X.size(), dtype = torch.float, device = device)
+        X = X.permute(1, 0) if batched else X
+        dX = dX.permute(1, 0) if batched else dX
+        dX[1:] = X[1:] - X[:-1] # Backward Difference
+        dX[0] = X[1] - X[0] # Forward Difference
+
+        dX = dX.permute(1, 0) if batched else dX
+        dX /= delta_lst[0]
+    elif dim == 2:
+        #print('dim = 2')
+        dX = torch.zeros(X.size() + tuple([2]), dtype = torch.float, device = device)
+        X = X.permute(1, 2, 0) if batched else X
+        dX = dX.permute(1, 2, 3, 0) if batched else dX # put batch to last dim
+        dX[1:, :, 0] = X[1:, :] - X[:-1, :] # Backward Difference
+        dX[0, :, 0] = X[1] - X[0] # Forward Difference
+
+        dX[:, 1:, 1] = X[:, 1:] - X[:, :-1] # Backward Difference
+        dX[:, 0, 1] = X[:, 1] - X[:, 0] # Forward Difference
+
+        dX = dX.permute(3, 0, 1, 2) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+    elif dim == 3:
+        #print('dim = 3')
+        dX = torch.zeros(X.size() + tuple([3]), dtype = torch.float, device = device)
+        X = X.permute(1, 2, 3, 0) if batched else X
+        dX = dX.permute(1, 2, 3, 4, 0) if batched else dX
+        dX[1:, :, :, 0] = X[1:, :, :] - X[:-1, :, :] # Backward Difference
+        dX[0, :, :, 0] = X[1] - X[0] # Forward Difference
+
+        dX[:, 1:, :, 1] = X[:, 1:] - X[:, :-1] # Backward Difference
+        dX[:, 0, :, 1] = X[:, 1] - X[:, 0] # Forward Difference
+
+        dX[:, :, 1:, 2] = X[:, :, 1:] - X[:, :, :-1] # Backward Difference
+        dX[:, :, 0, 2] = X[:, :, 1] - X[:, :, 0] # Forward Difference
+
+        dX = dX.permute(4, 0, 1, 2, 3) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+        dX[..., 2] /= delta_lst[2]
+    return dX
+  
+
+def gradient_c(X, batched = False, delta_lst = [1., 1., 1.]):
+    '''
+    Compute gradient of a torch tensor "X" in each direction
+    Non-boundaries: Central Difference
+    Upper-boundaries: Backward Difference
+    Lower-boundaries: Forward Difference
+    if X is batched: (n_batch, ...);
+    else: (...)
+    '''
+    device = X.device
+    dim = len(X.size()) - 1 if batched else len(X.size())
+    #print(X.size())
+    #print(batched)
+    #print(dim)
+    if dim == 1:
+        #print('dim = 1')
+        dX = torch.zeros(X.size(), dtype = torch.float, device = device)
+        X = X.permute(1, 0) if batched else X
+        dX = dX.permute(1, 0) if batched else dX
+        dX[1:-1] = (X[2:] - X[:-2]) / 2 # Central Difference
+        dX[0] = X[1] - X[0] # Forward Difference
+        dX[-1] = X[-1] - X[-2] # Backward Difference
+
+        dX = dX.permute(1, 0) if batched else dX
+        dX /= delta_lst[0]
+    elif dim == 2:
+        #print('dim = 2')
+        dX = torch.zeros(X.size() + tuple([2]), dtype = torch.float, device = device)
+        X = X.permute(1, 2, 0) if batched else X
+        dX = dX.permute(1, 2, 3, 0) if batched else dX # put batch to last dim
+        dX[1:-1, :, 0] = (X[2:, :] - X[:-2, :]) / 2
+        dX[0, :, 0] = X[1] - X[0]
+        dX[-1, :, 0] = X[-1] - X[-2]
+        dX[:, 1:-1, 1] = (X[:, 2:] - X[:, :-2]) / 2
+        dX[:, 0, 1] = X[:, 1] - X[:, 0]
+        dX[:, -1, 1] = X[:, -1] - X[:, -2]
+
+        dX = dX.permute(3, 0, 1, 2) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+    elif dim == 3:
+        #print('dim = 3')
+        dX = torch.zeros(X.size() + tuple([3]), dtype = torch.float, device = device)
+        X = X.permute(1, 2, 3, 0) if batched else X
+        dX = dX.permute(1, 2, 3, 4, 0) if batched else dX
+        dX[1:-1, :, :, 0] = (X[2:, :, :] - X[:-2, :, :]) / 2
+        dX[0, :, :, 0] = X[1] - X[0]
+        dX[-1, :, :, 0] = X[-1] - X[-2]
+        dX[:, 1:-1, :, 1] = (X[:, 2:, :] - X[:, :-2, :]) / 2
+        dX[:, 0, :, 1] = X[:, 1] - X[:, 0]
+        dX[:, -1, :, 1] = X[:, -1] - X[:, -2]
+        dX[:, :, 1:-1, 2] = (X[:, :, 2:] - X[:, :, :-2]) / 2
+        dX[:, :, 0, 2] = X[:, :, 1] - X[:, :, 0]
+        dX[:, :, -1, 2] = X[:, :, -1] - X[:, :, -2]
+
+        dX = dX.permute(4, 0, 1, 2, 3) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+        dX[..., 2] /= delta_lst[2]
+    return dX
+
+
+def gradient_c_numpy(X, batched = False, delta_lst = [1., 1., 1.]):
+    '''
+    Compute gradient of a Numpy array "X" in each direction
+    Non-boundaries: Central Difference
+    Upper-boundaries: Backward Difference
+    Lower-boundaries: Forward Difference
+    if X is batched: (n_batch, ...);
+    else: (...)
+    '''
+    dim = len(X.shape) - 1 if batched else len(X.shape)
+    #print(dim)
+    if dim == 1:
+        #print('dim = 1')
+        X = np.transpose(X, (1, 0)) if batched else X
+        dX = np.zeros(X.shapee).astype(float)
+        dX[1:-1] = (X[2:] - X[:-2]) / 2 # Central Difference
+        dX[0] = X[1] - X[0] # Forward Difference
+        dX[-1] = X[-1] - X[-2] # Backward Difference
+
+        dX = np.transpose(X, (1, 0)) if batched else dX
+        dX /= delta_lst[0]
+    elif dim == 2:
+        #print('dim = 2')
+        dX = np.zeros(X.shape + tuple([2])).astype(float)
+        X = np.transpose(X, (1, 2, 0)) if batched else X
+        dX = np.transpose(dX, (1, 2, 3, 0)) if batched else dX # put batch to last dim
+        dX[1:-1, :, 0] = (X[2:, :] - X[:-2, :]) / 2
+        dX[0, :, 0] = X[1] - X[0]
+        dX[-1, :, 0] = X[-1] - X[-2]
+        dX[:, 1:-1, 1] = (X[:, 2:] - X[:, :-2]) / 2
+        dX[:, 0, 1] = X[:, 1] - X[:, 0]
+        dX[:, -1, 1] = X[:, -1] - X[:, -2]
+
+        dX = np.transpose(dX, (3, 0, 1, 2)) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+    elif dim == 3:
+        #print('dim = 3')
+        dX = np.zeros(X.shape + tuple([3])).astype(float)
+        X = np.transpose(X, (1, 2, 3, 0)) if batched else X
+        dX = np.transpose(dX, (1, 2, 3, 4, 0)) if batched else dX # put batch to last dim
+        dX[1:-1, :, :, 0] = (X[2:, :, :] - X[:-2, :, :]) / 2
+        dX[0, :, :, 0] = X[1] - X[0]
+        dX[-1, :, :, 0] = X[-1] - X[-2]
+        dX[:, 1:-1, :, 1] = (X[:, 2:, :] - X[:, :-2, :]) / 2
+        dX[:, 0, :, 1] = X[:, 1] - X[:, 0]
+        dX[:, -1, :, 1] = X[:, -1] - X[:, -2]
+        dX[:, :, 1:-1, 2] = (X[:, :, 2:] - X[:, :, :-2]) / 2
+        dX[:, :, 0, 2] = X[:, :, 1] - X[:, :, 0]
+        dX[:, :, -1, 2] = X[:, :, -1] - X[:, :, -2]
+
+        dX = np.transpose(dX, (4, 0, 1, 2, 3)) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+        dX[..., 2] /= delta_lst[2]
+    return dX
+
+
+def gradient_f_numpy(X, batched = False, delta_lst = [1., 1., 1.]):
+    '''
+    Compute gradient of a torch tensor "X" in each direction
+    Upper-boundaries: Backward Difference
+    Non-boundaries & Upper-boundaries: Forward Difference
+    if X is batched: (n_batch, ...);
+    else: (...)
+    '''
+    dim = len(X.shape) - 1 if batched else len(X.shape)
+    #print(dim)
+    if dim == 1:
+        #print('dim = 1')
+        X = np.transpose(X, (1, 0)) if batched else X
+        dX = np.zeros(X.shapee).astype(float)
+        dX[-1] = X[-1] - X[-2] # Backward Difference
+        dX[:-1] = X[1:] - X[:-1] # Forward Difference
+
+        dX = np.transpose(X, (1, 0)) if batched else dX
+        dX /= delta_lst[0]
+    elif dim == 2:
+        #print('dim = 2')
+        dX = np.zeros(X.shape + tuple([2])).astype(float)
+        X = np.transpose(X, (1, 2, 0)) if batched else X
+        dX = np.transpose(dX, (1, 2, 3, 0)) if batched else dX # put batch to last dim
+        dX[-1, :, 0] = X[-1, :] - X[-2, :] # Backward Difference
+        dX[:-1, :, 0] = X[1:] - X[:-1] # Forward Difference
+
+        dX[:, -1, 1] = X[:, -1] - X[:, -2] # Backward Difference
+        dX[:, :-1, 1] = X[:, 1:] - X[:, :-1] # Forward Difference
+
+        dX = np.transpose(dX, (3, 0, 1, 2)) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+    elif dim == 3:
+        #print('dim = 3')
+        dX = np.zeros(X.shape + tuple([3])).astype(float)
+        X = np.transpose(X, (1, 2, 3, 0)) if batched else X
+        dX = np.transpose(dX, (1, 2, 3, 4, 0)) if batched else dX # put batch to last dim
+        dX[-1, :, :, 0] = X[-1, :, :] - X[-2, :, :] # Backward Difference
+        dX[:-1, :, :, 0] = X[1:] - X[:-1] # Forward Difference
+
+        dX[:, -1, :, 1] = X[:, -1] - X[:, -2] # Backward Difference
+        dX[:, :-1, :, 1] = X[:, 1:] - X[:, :-1] # Forward Difference
+
+        dX[:, :, -1, 2] = X[:, :, -1] - X[:, :, -2] # Backward Difference
+        dX[:, :, :-1, 2] = X[:, :, 1:] - X[:, :, :-1] # Forward Difference
+
+        dX = np.transpose(dX, (4, 0, 1, 2, 3)) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+        dX[..., 2] /= delta_lst[2]
+    return dX
+
+
+class Upwind(object):
+    '''
+    Backward if > 0, forward if <= 0
+    '''
+    def __init__(self, U, data_spacing = [1., 1, 1.], batched = True):
+        self.U = U # (s, r, c)
+        self.batched = batched
+        self.data_spacing = data_spacing
+        self.dim = len(self.U.size()) - 1 if batched else len(self.U.size())
+        self.I = torch.ones(self.U.size(), dtype = torch.float, device = U.device)
+
+    def dX(self, FGx):
+        dXf = gradient_f(self.U, batched = self.batched, delta_lst = self.data_spacing)[..., 0]
+        dXb = gradient_b(self.U, batched = self.batched, delta_lst = self.data_spacing)[..., 0]
+        Xflag = (FGx > 0).float()
+        return dXf * (self.I - Xflag) + dXb * Xflag
+
+    def dY(self, FGy):
+        dYf = gradient_f(self.U, batched = self.batched, delta_lst = self.data_spacing)[..., 1]
+        dYb = gradient_b(self.U, batched = self.batched, delta_lst = self.data_spacing)[..., 1]
+        Yflag = (FGy > 0).float()
+        return dYf * (self.I - Yflag) + dYb * Yflag
+
+    def dZ(self, FGz):
+        dZf = gradient_f(self.U, batched = self.batched, delta_lst = self.data_spacing)[..., 2]
+        dZb = gradient_b(self.U, batched = self.batched, delta_lst = self.data_spacing)[..., 2]
+        Zflag = (FGz > 0).float()
+        return dZf * (self.I - Zflag) + dZb * Zflag
+    
+    
+class AdvDiffPartial(nn.Module):
+    def __init__(self, data_spacing, device):
+        super(AdvDiffPartial, self).__init__()
+        self.dimension = len(data_spacing)  # (slc, row, col)
+        self.device = device
+        self.data_spacing = data_spacing
+
+    @property
+    def Grad_Ds(self):
+        return {
+            'constant': self.Grad_constantD,
+            'scalar': self.Grad_scalarD,
+            'diag': self.Grad_diagD,
+            'full': self.Grad_fullD,
+            'full_dual': self.Grad_fullD,
+            'full_spectral':self.Grad_fullD,
+            'full_cholesky': self.Grad_fullD,
+            'full_symmetric': self.Grad_fullD
+        }
+    @property
+    def Grad_Vs(self):
+        return {
+            'constant': self.Grad_constantV,
+            'scalar': self.Grad_scalarV,
+            'vector': self.Grad_vectorV, # For general V w/o div-free TODO self.Grad_vectorV
+            'vector_div_free': self.Grad_div_free_vectorV,
+            'vector_div_free_clebsch': self.Grad_div_free_vectorV,
+            'vector_div_free_stream': self.Grad_div_free_vectorV,
+            'vector_div_free_stream_gauge': self.Grad_div_free_vectorV, 
+        }
+
+    def Grad_constantD(self, C, Dlst):
+        if self.dimension == 1:
+            return Dlst['D'] * (self.ddXc(C))
+        elif self.dimension == 2:
+            return Dlst['D'] * (self.ddXc(C) + self.ddYc(C))
+        elif self.dimension == 3:
+            return Dlst['D'] * (self.ddXc(C) + self.ddYc(C) + self.ddZc(C))
+
+    def Grad_constant_tensorD(self, C, Dlst):
+        if self.dimension == 1:
+            raise NotImplementedError
+        elif self.dimension == 2:
+            dC_c = self.dc(C)
+            dC_f = self.df(C)
+            return Dlst['Dxx'] * self.dXb(dC_f[..., 0]) +\
+                Dlst['Dxy'] * self.dXb(dC_f[..., 1]) + Dlst['Dxy'] * self.dYb(dC_f[..., 0]) +\
+                        Dlst['Dyy'] * self.dYb(dC_f[..., 1])  
+        elif self.dimension == 3:
+            dC_c = self.dc(C)
+            dC_f = self.df(C)
+            return Dlst['Dxx'] * self.dXb(dC_f[..., 0]) + Dlst['Dyy'] * self.dYb(dC_f[..., 1]) + Dlst['Dzz'] * self.dZb(dC_f[..., 2]) + \
+                            Dlst['Dxy'] * (self.dXb(dC_f[..., 1]) + self.dYb(dC_f[..., 0])) + \
+                                Dlst['Dyz'] * (self.dYb(dC_f[..., 2]) + self.dZb(dC_f[..., 1])) + \
+                                    Dlst['Dxz'] * (self.dZb(dC_f[..., 0]) + self.dXb(dC_f[..., 2]))
+        
+    def Grad_scalarD(self, C, Dlst): # batch_C: (batch_size, (slc), row, col)
+        # Expanded version: \nabla (D \nabla C) => \nabla D \cdot \nabla C (part (a)) + D \Delta C (part (b)) # 
+        # NOTE: Work better than Central Differences !!! # 
+        # Nested Forward-Backward Difference Scheme in part (b)#
+        if self.dimension == 1:
+            dC = gradient_c(C, batched = True, delta_lst = self.data_spacing)
+            return gradient_c(Dlst['D'], batched = True, delta_lst = self.data_spacing) * dC + \
+                Dlst['D'] * gradient_c(dC, batched = True, delta_lst = self.data_spacing)
+        else: # (dimension = 2 or 3)
+            dC_c = gradient_c(C, batched = True, delta_lst = self.data_spacing)
+            dC_f = gradient_f(C, batched = True, delta_lst = self.data_spacing)
+            dD_c = gradient_c(Dlst['D'], batched = True, delta_lst = self.data_spacing)
+            out = (dD_c * dC_c).sum(-1)
+            for dim in range(dC_f.size(-1)):
+                out += Dlst['D'] * gradient_b(dC_f[..., dim], batched = True, delta_lst = self.data_spacing)[..., dim]
+            return out
+
+    def Grad_diagD(self, C, Dlst):
+        # Expanded version #
+        if self.dimension == 1:
+            raise NotImplementedError('diag_D is not supported for 1D version of diffusivity')
+        elif self.dimension == 2:
+            dC_c = self.dc(C)
+            dC_f = self.df(C)
+            return self.dXc(Dlst['Dxx']) * dC_c[..., 0] + Dlst['Dxx'] * self.dXb(dC_f[..., 0]) +\
+                self.dYc(Dlst['Dyy']) * dC_c[..., 1] + Dlst['Dyy'] * self.dYb(dC_f[..., 1]) 
+        elif self.dimension == 3:
+            dC_c = self.dc(C)
+            dC_f = self.df(C)
+            return self.dXc(Dlst['Dxx']) * dC_c[..., 0] + Dlst['Dxx'] * self.dXb(dC_f[..., 0]) +\
+                self.dYc(Dlst['Dyy']) * dC_c[..., 1] + Dlst['Dyy'] * self.dYb(dC_f[..., 1]) +\
+                    self.dZc(Dlst['Dzz']) * dC_c[..., 2] + Dlst['Dzz'] * self.dZb(dC_f[..., 2])
+
+    def Grad_fullD(self, C, Dlst):
+        # Expanded version #
+        '''https://github.com/uncbiag/PIANOinD/blob/master/Doc/PIANOinD.pdf'''
+        if self.dimension == 1:
+            raise NotImplementedError('full_D is not supported for 1D version of diffusivity')
+        elif self.dimension == 2:
+            dC_c = self.dc(C)
+            dC_f = self.df(C)
+            return self.dXc(Dlst['Dxx']) * dC_c[..., 0] + Dlst['Dxx'] * self.dXb(dC_f[..., 0]) +\
+                self.dXc(Dlst['Dxy']) * dC_c[..., 1] + Dlst['Dxy'] * self.dXb(dC_f[..., 1]) +\
+                    self.dYc(Dlst['Dxy']) * dC_c[..., 0] + Dlst['Dxy'] * self.dYb(dC_f[..., 0]) +\
+                        self.dYc(Dlst['Dyy']) * dC_c[..., 1] + Dlst['Dyy'] * self.dYb(dC_f[..., 1])  
+        elif self.dimension == 3:
+            dC_c = self.dc(C)
+            dC_f = self.df(C)
+            return (self.dXc(Dlst['Dxx']) + self.dYc(Dlst['Dxy']) + self.dZc(Dlst['Dxz'])) * dC_c[..., 0] + \
+                (self.dXc(Dlst['Dxy']) + self.dYc(Dlst['Dyy']) + self.dZc(Dlst['Dyz'])) * dC_c[..., 1] + \
+                    (self.dXc(Dlst['Dxz']) + self.dYc(Dlst['Dyz']) + self.dZc(Dlst['Dzz'])) * dC_c[..., 2] + \
+                        Dlst['Dxx'] * self.dXb(dC_f[..., 0]) + Dlst['Dyy'] * self.dYb(dC_f[..., 1]) + Dlst['Dzz'] * self.dZb(dC_f[..., 2]) + \
+                            Dlst['Dxy'] * (self.dXb(dC_f[..., 1]) + self.dYb(dC_f[..., 0])) + \
+                                Dlst['Dyz'] * (self.dYb(dC_f[..., 2]) + self.dZb(dC_f[..., 1])) + \
+                                    Dlst['Dxz'] * (self.dZb(dC_f[..., 0]) + self.dXb(dC_f[..., 2]))
+
+    def Grad_constantV(self, C, Vlst):
+        if len(Vlst['V'].size()) == 1:
+            if self.dimension == 1:
+                return - Vlst['V'] * self.dXb(C) if Vlst['V'] > 0 else - Vlst['V'] * self.dXf(C)
+            elif self.dimension == 2:
+                return - Vlst['V'] * (self.dXb(C) + self.dYb(C)) if Vlst['V'] > 0 else - Vlst['V'] * (self.dXf(C) + self.dYf(C))
+            elif self.dimension == 3:
+                return - Vlst['V'] * (self.dXb(C) + self.dYb(C) + self.dZb(C)) if Vlst['V'] > 0 else - Vlst['V'] * (self.dXf(C) + self.dYf(C) + self.dZf(C))
+        else:
+            if self.dimension == 1:
+                return - Vlst['V'] * self.dXb(C) if Vlst['V'][0, 0] > 0 else - Vlst['V'] * self.dXf(C)
+            elif self.dimension == 2:
+                return - Vlst['V'] * (self.dXb(C) + self.dYb(C)) if Vlst['V'][0, 0, 0] > 0 else - Vlst['V'] * (self.dXf(C) + self.dYf(C))
+            elif self.dimension == 3:
+                return - Vlst['V'] * (self.dXb(C) + self.dYb(C) + self.dZb(C)) if Vlst['V'][0, 0, 0, 0] > 0 else - Vlst['V'] * (self.dXf(C) + self.dYf(C) + self.dZf(C))
+    
+    def Grad_constant_vectorV(self, C, Vlst):
+        if self.dimension == 1:
+            raise NotImplementedError
+        elif self.dimension == 2:
+            out_x = - Vlst['Vx'] * (self.dXb(C) + self.dYb(C)) if Vlst['Vx'][0, 0, 0] > 0 else - Vlst['Vx'] * (self.dXf(C) + self.dYf(C))
+            out_y = - Vlst['Vy'] * (self.dXb(C) + self.dYb(C)) if Vlst['Vy'][0, 0, 0] > 0 else - Vlst['Vy'] * (self.dXf(C) + self.dYf(C))
+            return out_x + out_y
+        elif self.dimension == 3:
+            out_x = - Vlst['Vx'] * (self.dXb(C) + self.dYb(C)) if Vlst['Vx'][0, 0, 0] > 0 else - Vlst['Vx'] * (self.dXf(C) + self.dYf(C))
+            out_y = - Vlst['Vy'] * (self.dXb(C) + self.dYb(C)) if Vlst['Vy'][0, 0, 0] > 0 else - Vlst['Vy'] * (self.dXf(C) + self.dYf(C))
+            out_z = - Vlst['Vz'] * (self.dXb(C) + self.dYb(C)) if Vlst['Vz'][0, 0, 0] > 0 else - Vlst['Vz'] * (self.dXf(C) + self.dYf(C))
+            return out_x + out_y + out_z
+    
+    def Grad_SimscalarV(self, C, Vlst):
+        V = Vlst['V']
+        Upwind_C = Upwind(C, self.data_spacing)
+        if self.dimension == 1:
+            C_x = Upwind_C.dX(V)
+            return - V * C_x
+        if self.dimension == 2:
+            C_x, C_y = Upwind_C.dX(V), Upwind_C.dY(V)
+            return - V * (C_x + C_y)
+        if self.dimension == 3:
+            C_x, C_y, C_z = Upwind_C.dX(V), Upwind_C.dY(V), Upwind_C.dZ(V)
+            return - V * (C_x + C_y + C_z)
+
+    def Grad_scalarV(self, C, Vlst):
+        V = Vlst['V']
+        Upwind_C = Upwind(C, self.data_spacing)
+        dV = gradient_c(V, batched = True, delta_lst = self.data_spacing)
+        if self.dimension == 1:
+            C_x = Upwind_C.dX(V)
+            return - V * C_x - C * dV
+        elif self.dimension == 2:
+            C_x, C_y = Upwind_C.dX(V), Upwind_C.dY(V)
+            return - V * (C_x + C_y) - C * dV.sum(-1)
+        elif self.dimension == 3:
+            C_x, C_y, C_z = Upwind_C.dX(V), Upwind_C.dY(V), Upwind_C.dZ(V)
+            return - V * (C_x + C_y + C_z) - C * dV.sum(-1)
+
+    def Grad_div_free_vectorV(self, C, Vlst):
+        ''' For divergence-free-by-definition velocity'''
+        if self.dimension == 1:
+            raise NotImplementedError('clebschVector is not supported for 1D version of velocity')
+        Upwind_C = Upwind(C, self.data_spacing) 
+        C_x, C_y = Upwind_C.dX(Vlst['Vx']), Upwind_C.dY(Vlst['Vy'])
+        if self.dimension == 2:
+            return - (Vlst['Vx'] * C_x + Vlst['Vy'] * C_y)
+        elif self.dimension == 3:
+            C_z = Upwind_C.dZ(Vlst['Vz'])
+            return - (Vlst['Vx'] * C_x + Vlst['Vy'] * C_y + Vlst['Vz'] * C_z)
+            
+    def Grad_vectorV(self, C, Vlst):
+        ''' For general velocity'''
+        if self.dimension == 1:
+            raise NotImplementedError('vector is not supported for 1D version of velocity')
+        Upwind_C = Upwind(C, self.data_spacing) 
+        C_x, C_y = Upwind_C.dX(Vlst['Vx']), Upwind_C.dY(Vlst['Vy'])
+        Vx_x = self.dXc(Vlst['Vx'])
+        Vy_y = self.dYc(Vlst['Vy']) 
+        if self.dimension == 2:
+            return - (Vlst['Vx'] * C_x + Vlst['Vy'] * C_y) - C * (Vx_x + Vy_y)
+        if self.dimension == 3:
+            C_z = Upwind_C.dZ(Vlst['Vz'])
+            Vz_z = self.dZc(Vlst['Vz'])
+            return - (Vlst['Vx'] * C_x + Vlst['Vy'] * C_y + Vlst['Vz'] * C_z) - C * (Vx_x + Vy_y + Vz_z)
+    
+    ################# Utilities #################
+    def db(self, X):
+        return gradient_b(X, batched = True, delta_lst = self.data_spacing)
+    def df(self, X):
+        return gradient_f(X, batched = True, delta_lst = self.data_spacing)
+    def dc(self, X):
+        return gradient_c(X, batched = True, delta_lst = self.data_spacing)
+    def dXb(self, X):
+        return gradient_b(X, batched = True, delta_lst = self.data_spacing)[..., 0]
+    def dXf(self, X):
+        return gradient_f(X, batched = True, delta_lst = self.data_spacing)[..., 0]
+    def dXc(self, X):
+        return gradient_c(X, batched = True, delta_lst = self.data_spacing)[..., 0]
+    def dYb(self, X):
+        return gradient_b(X, batched = True, delta_lst = self.data_spacing)[..., 1]
+    def dYf(self, X):
+        return gradient_f(X, batched = True, delta_lst = self.data_spacing)[..., 1]
+    def dYc(self, X):
+        return gradient_c(X, batched = True, delta_lst = self.data_spacing)[..., 1]
+    def dZb(self, X):
+        return gradient_b(X, batched = True, delta_lst = self.data_spacing)[..., 2]
+    def dZf(self, X):
+        return gradient_f(X, batched = True, delta_lst = self.data_spacing)[..., 2]
+    def dZc(self, X):
+        return gradient_c(X, batched = True, delta_lst = self.data_spacing)[..., 2]
+    def ddXc(self, X):
+        return gradient_b(gradient_f(X, batched = True, delta_lst = self.data_spacing)[..., 0], 
+        batched = True, delta_lst = self.data_spacing)[..., 0]
+    def ddYc(self, X):
+        return gradient_b(gradient_f(X, batched = True, delta_lst = self.data_spacing)[..., 1], 
+        batched = True, delta_lst = self.data_spacing)[..., 1]
+    def ddZc(self, X):
+        return gradient_b(gradient_f(X, batched = True, delta_lst = self.data_spacing)[..., 2], 
+        batched = True, delta_lst = self.data_spacing)[..., 2]
+    
+
+
+class AdvDiffPDE(nn.Module):
+    '''
+    Plain advection-diffusion PDE solver for pre-set V_lst and D_lst (1D, 2D, 3D) for forward time series simulation
+    '''
+    def __init__(self, data_spacing, perf_pattern, D_type='scalar', V_type='vector', BC=None, dt=0.1, V_dict={}, D_dict={}, stochastic=False, device='cpu'):
+        super(AdvDiffPDE, self).__init__() 
+        self.BC = BC
+        self.dt = dt
+        self.dimension = len(data_spacing)
+        self.perf_pattern = perf_pattern
+        self.partials = AdvDiffPartial(data_spacing, device)
+        self.D_type, self.V_type = D_type, V_type
+        self.stochastic = stochastic 
+        self.V_dict, self.D_dict = V_dict, D_dict
+        self.Sigma, self.Sigma_V, self.Sigma_D = 0., 0., 0. # Only for initialization # 
+        if self.dimension == 1:
+            self.neumann_BC = torch.nn.ReplicationPad1d(1)
+        elif self.dimension == 2:
+            self.neumann_BC = torch.nn.ReplicationPad2d(1)
+        elif self.dimension == 3:
+            self.neumann_BC = torch.nn.ReplicationPad3d(1)
+        else:
+            raise ValueError('Unsupported dimension: %d' % self.dimension)
+                 
+    @property
+    def set_BC(self):
+    # NOTE For bondary condition of mass concentration #
+        '''X: (n_batch, spatial_shape)'''
+        if self.BC == 'neumann' or self.BC == 'cauchy':
+            if self.dimension == 1:
+                return lambda X: self.neumann_BC(X[:, 1:-1].unsqueeze(dim=1))[:,0]
+            elif self.dimension == 2:
+                return lambda X: self.neumann_BC(X[:, 1:-1, 1:-1].unsqueeze(dim=1))[:,0]
+            elif self.dimension == 3:
+                return lambda X: self.neumann_BC(X[:, 1:-1, 1:-1, 1:-1].unsqueeze(dim=1))[:,0]
+            else:
+                raise NotImplementedError('Unsupported B.C.!')
+        elif self.BC == 'dirichlet_neumann' or self.BC == 'source_neumann':
+            ctrl_wdth = 1
+            if self.dimension == 1:
+                self.dirichlet_BC = torch.nn.ReplicationPad1d(ctrl_wdth)
+                return lambda X: self.dirichlet_BC(X[:, ctrl_wdth : -ctrl_wdth].unsqueeze(dim=1))[:,0]
+            elif self.dimension == 2:
+                self.dirichlet_BC = torch.nn.ReplicationPad2d(ctrl_wdth)
+                return lambda X: self.dirichlet_BC(X[:, ctrl_wdth : -ctrl_wdth, ctrl_wdth : -ctrl_wdth].unsqueeze(dim=1))[:,0]
+            elif self.dimension == 3:
+                self.dirichlet_BC = torch.nn.ReplicationPad3d(ctrl_wdth)
+                return lambda X: self.neumann_dirichlet_BCBC(X[:, ctrl_wdth : -ctrl_wdth, ctrl_wdth : -ctrl_wdth, ctrl_wdth : -ctrl_wdth].unsqueeze(dim=1))[:,0]
+            else:
+                raise NotImplementedError('Unsupported B.C.!')
+        else:
+            return lambda X: X 
+        
+    def forward(self, t, batch_C):
+        '''
+        t: (batch_size,)
+        batch_C: (batch_size, (slc,) row, col)
+        ''' 
+        batch_size  = batch_C.size(0)
+        batch_C = self.set_BC(batch_C)
+        if 'diff' not in self.perf_pattern:
+            out = self.partials.Grad_Vs[self.V_type](batch_C, self.V_dict) 
+            if self.stochastic:  
+                out = out + self.Sigma * math.sqrt(self.dt) * torch.randn_like(batch_C).to(batch_C)
+        elif 'adv' not in self.perf_pattern:
+            out = self.partials.Grad_Ds[self.D_type](batch_C, self.D_dict)
+            if self.stochastic:  
+                out = out + self.Sigma * math.sqrt(self.dt) * torch.randn_like(batch_C).to(batch_C)
+        else:
+            if self.stochastic:  
+                out_D = self.partials.Grad_Ds[self.D_type](batch_C, self.D_dict)
+                out_V = self.partials.Grad_Vs[self.V_type](batch_C, self.V_dict) 
+                out = out_D + out_V + self.Sigma * math.sqrt(self.dt) * torch.randn_like(batch_C).to(batch_C) 
+            else:
+                out_V = self.partials.Grad_Vs[self.V_type](batch_C, self.V_dict)  
+                out_D = self.partials.Grad_Ds[self.D_type](batch_C, self.D_dict)
+                out = out_V + out_D
+        return out
+
+        
+

ShapeID/DiffEqs/rk_common.py

@@ -0,0 +1,78 @@
+# Based on https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/integrate
+import collections
+from ShapeID.DiffEqs.misc import _scaled_dot_product, _convert_to_tensor
+
+_ButcherTableau = collections.namedtuple('_ButcherTableau', 'alpha beta c_sol c_error')
+
+
+class _RungeKuttaState(collections.namedtuple('_RungeKuttaState', 'y1, f1, t0, t1, dt, interp_coeff')):
+    """Saved state of the Runge Kutta solver.
+
+    Attributes:
+        y1: Tensor giving the function value at the end of the last time step.
+        f1: Tensor giving derivative at the end of the last time step.
+        t0: scalar float64 Tensor giving start of the last time step.
+        t1: scalar float64 Tensor giving end of the last time step.
+        dt: scalar float64 Tensor giving the size for the next time step.
+        interp_coef: list of Tensors giving coefficients for polynomial
+            interpolation between `t0` and `t1`.
+    """
+
+
+def _runge_kutta_step(func, y0, f0, t0, dt, tableau):
+    """Take an arbitrary Runge-Kutta step and estimate error.
+
+    Args:
+        func: Function to evaluate like `func(t, y)` to compute the time derivative
+            of `y`.
+        y0: Tensor initial value for the state.
+        f0: Tensor initial value for the derivative, computed from `func(t0, y0)`.
+        t0: float64 scalar Tensor giving the initial time.
+        dt: float64 scalar Tensor giving the size of the desired time step.
+        tableau: optional _ButcherTableau describing how to take the Runge-Kutta
+            step.
+        name: optional name for the operation.
+
+    Returns:
+        Tuple `(y1, f1, y1_error, k)` giving the estimated function value after
+        the Runge-Kutta step at `t1 = t0 + dt`, the derivative of the state at `t1`,
+        estimated error at `t1`, and a list of Runge-Kutta coefficients `k` used for
+        calculating these terms.
+    """
+    dtype = y0[0].dtype
+    device = y0[0].device
+
+    t0 = _convert_to_tensor(t0, dtype=dtype, device=device)
+    dt = _convert_to_tensor(dt, dtype=dtype, device=device)
+
+    k = tuple(map(lambda x: [x], f0))
+    for alpha_i, beta_i in zip(tableau.alpha, tableau.beta):
+        ti = t0 + alpha_i * dt
+        yi = tuple(y0_ + _scaled_dot_product(dt, beta_i, k_) for y0_, k_ in zip(y0, k))
+        tuple(k_.append(f_) for k_, f_ in zip(k, func(ti, yi)))
+
+    if not (tableau.c_sol[-1] == 0 and tableau.c_sol[:-1] == tableau.beta[-1]):
+        # This property (true for Dormand-Prince) lets us save a few FLOPs.
+        yi = tuple(y0_ + _scaled_dot_product(dt, tableau.c_sol, k_) for y0_, k_ in zip(y0, k))
+
+    y1 = yi
+    f1 = tuple(k_[-1] for k_ in k)
+    y1_error = tuple(_scaled_dot_product(dt, tableau.c_error, k_) for k_ in k)
+    return (y1, f1, y1_error, k)
+
+
+def rk4_step_func(func, t, dt, y, k1=None):
+    if k1 is None: k1 = func(t, y)
+    k2 = func(t + dt / 2, tuple(y_ + dt * k1_ / 2 for y_, k1_ in zip(y, k1)))
+    k3 = func(t + dt / 2, tuple(y_ + dt * k2_ / 2 for y_, k2_ in zip(y, k2)))
+    k4 = func(t + dt, tuple(y_ + dt * k3_ for y_, k3_ in zip(y, k3)))
+    return tuple((k1_ + 2 * k2_ + 2 * k3_ + k4_) * (dt / 6) for k1_, k2_, k3_, k4_ in zip(k1, k2, k3, k4))
+
+
+def rk4_alt_step_func(func, t, dt, y, k1=None):
+    """Smaller error with slightly more compute."""
+    if k1 is None: k1 = func(t, y)
+    k2 = func(t + dt / 3, tuple(y_ + dt * k1_ / 3 for y_, k1_ in zip(y, k1)))
+    k3 = func(t + dt * 2 / 3, tuple(y_ + dt * (k1_ / -3 + k2_) for y_, k1_, k2_ in zip(y, k1, k2)))
+    k4 = func(t + dt, tuple(y_ + dt * (k1_ - k2_ + k3_) for y_, k1_, k2_, k3_ in zip(y, k1, k2, k3)))
+    return tuple((k1_ + 3 * k2_ + 3 * k3_ + k4_) * (dt / 8) for k1_, k2_, k3_, k4_ in zip(k1, k2, k3, k4))

ShapeID/DiffEqs/solvers.py

@@ -0,0 +1,216 @@
+import abc
+import torch
+from ShapeID.DiffEqs.misc import _assert_increasing, _handle_unused_kwargs
+
+def set_BC_2D(X, BCs): # X: (n_batch, spatial_size); BCs: (batch, 4, BC_shape, data_dim)
+    BC_size = BCs.size(2)
+    X[:, : BC_size] = BCs[:, 0]
+    X[:, - BC_size :] = BCs[:, 1]
+    X[:, :, : BC_size] = BCs[:, 2].permute(0, 2, 1) # (batch, BC_shape, r) -> (batch, r, BC_shape)
+    X[:, :, - BC_size :] = BCs[:, 3].permute(0, 2, 1) # (batch, BC_shape, r) -> (batch, r, BC_shape)
+    del BCs
+    return X
+def set_BC_3D(X, BCs): # X: (n_batch, spatial_size); BCs: (batch, 6, BC_shape, data_dim, dta_dim)
+    BC_size = BCs.size(2)
+    X[:, : BC_size] = BCs[:, 0]
+    X[:, - BC_size :] = BCs[:, 1]
+    X[:, :, : BC_size] = BCs[:, 2].permute(0, 2, 1, 3) # (batch, BC_shape, s, c) -> (batch, s, BC_shape, c)
+    X[:, :, - BC_size :] = BCs[:, 3].permute(0, 2, 1, 3) # (batch, BC_shape, s, c) -> (batch, s, BC_shape, c)
+    X[:, :, :, : BC_size] = BCs[:, 4].permute(0, 2, 3, 1) # (batch, BC_shape, s, r) -> (batch, s, r, BC_shape)
+    X[:, :, :, - BC_size :] = BCs[:, 5].permute(0, 2, 3, 1) # (batch, BC_shape, s, r) -> (batch, s, r, BC_shape)
+    del BCs
+    return X
+
+''' X[t] = X[t] + dBC[t] (dBC[t] = BC[t+1] - BC[t]) '''
+def add_dBC_2D(X, dBCs): # X: (n_batch, spatial_size); BCs: (batch, 4, BC_shape, data_dim)
+    BC_size = dBCs.size(2)
+    X[:, : BC_size] += dBCs[:, 0]
+    X[:, - BC_size :] += dBCs[:, 1]
+    X[:, :, : BC_size] += dBCs[:, 2].permute(0, 2, 1) # (batch, BC_shape, r) -> (batch, r, BC_shape)
+    X[:, :, - BC_size :] += dBCs[:, 3].permute(0, 2, 1) # (batch, BC_shape, r) -> (batch, r, BC_shape)
+    del dBCs
+    return X
+def add_dBC_3D(X, dBCs): # X: (n_batch, spatial_size); BCs: (batch, 6, BC_shape, data_dim, dta_dim)
+    BC_size = dBCs.size(2)
+    X[:, : BC_size] += dBCs[:, 0]
+    X[:, - BC_size :] += dBCs[:, 1]
+    X[:, :, : BC_size] += dBCs[:, 2].permute(0, 2, 1, 3) # (batch, BC_shape, s, c) -> (batch, s, BC_shape, c)
+    X[:, :, - BC_size :] += dBCs[:, 3].permute(0, 2, 1, 3) # (batch, BC_shape, s, c) -> (batch, s, BC_shape, c)
+    X[:, :, :, : BC_size] += dBCs[:, 4].permute(0, 2, 3, 1) # (batch, BC_shape, s, r) -> (batch, s, r, BC_shape)
+    X[:, :, :, - BC_size :] += dBCs[:, 5].permute(0, 2, 3, 1) # (batch, BC_shape, s, r) -> (batch, s, r, BC_shape)
+    del dBCs
+    return X
+
+class AdaptiveStepsizeODESolver(object):
+    __metaclass__ = abc.ABCMeta
+
+    def __init__(self, func, y0, atol, rtol, options= None):
+        
+       # _handle_unused_kwargs(self, options)
+        #del options
+        self.func = func
+        self.y0 = y0
+        self.atol = atol
+        self.rtol = rtol
+
+    def before_integrate(self, t):
+        pass
+
+    @abc.abstractmethod
+    def advance(self, next_t):
+        raise NotImplementedError
+
+    def integrate(self, t):
+        _assert_increasing(t)
+        solution = [self.y0]
+        t = t.to(self.y0[0].device, torch.float64)
+        self.before_integrate(t)
+        for i in range(1, len(t)):
+            y = self.advance(t[i])
+            solution.append(y)
+        '''if self.contours is not None: # contours: (n_batch, nT, 4 / 6, BC_size, c)
+            if self.adjoint:
+                for i in range(1, len(t)):
+                    ys = list(self.advance(t[i])) # tuple: (y0, **back_grad) -> y0: (n_batch, spatial_shape)
+                    #print(len(t))
+                    #print(ys[0].size())
+                    #print(self.contours.size())
+                    ys[0] = self.set_BC(ys[0], self.contours[:, i]) # (n_batch, 4 / 6, BC_size, c)
+                    solution.append(tuple(ys))
+            else:
+                for i in range(1, len(t)):
+                    y = torch.stack(self.advance(t[i])) # y: (n_batch, 1, spatial_shape) 
+                    y = self.set_BC(y[:, 0], self.contours[:, i]).unsqueeze(1) 
+                    solution.append(tuple(y))
+        elif self.dcontours is not None: # dcontours: (n_batch, nT, 4 / 6, BC_size, c)
+            if self.adjoint:
+                for i in range(1, len(t)):
+                    ys = list(self.advance(t[i])) # ys - tuple: (y0, **back_grad) -> y0: (n_batch, spatial_shape)
+                    ys[0] = self.add_dBC(ys[0], self.dcontours[:, i]) # (n_batch, 4 / 6, BC_size, c)
+                    solution.append(tuple(ys))
+            else:
+                for i in range(1, len(t)):
+                    y = torch.stack(self.advance(t[i])) # (n_batch, 1, spatial_shape)
+                    y = self.add_dBC(y[:, 0], self.dcontours[:, i]).unsqueeze(1)
+                    solution.append(tuple(y))
+        else:
+            for i in range(1, len(t)):
+                y = self.advance(t[i])
+                solution.append(y)'''
+        return tuple(map(torch.stack, tuple(zip(*solution))))
+
+
+class FixedGridODESolver(object):
+    __metaclass__ = abc.ABCMeta
+
+    def __init__(self, func, y0, step_size=None, grid_constructor=None, atol=None, rtol=None, dt=None, options = None):
+        '''if 'dirichlet' in options.BC or 'cauchy' in options.BC and options.contours is not None:
+            self.contours = options.contours  # (n_batch, nT, 4 / 6, BC_size, sub_spatial_shape)
+            self.BC_size = self.contours.size(3)
+            self.set_BC = set_BC_2D if self.contours.size(2) == 4 else set_BC_3D
+        else:
+            self.contours = None
+        if 'source' in options.BC and options.dcontours is not None:
+            self.dcontours = options.dcontours  # (n_batch, nT, 4 / 6, BC_size, sub_spatial_shape)
+            self.BC_size = self.dcontours.size(3)
+            self.add_dBC = add_dBC_2D if self.dcontours.size(2) == 4 else add_dBC_3D
+        else:
+            self.dcontours = None'''
+        #self.adjoint = options.adjoint
+        #options.pop('rtol', None)
+        #options.pop('atol', None)
+        #_handle_unused_kwargs(self, options)
+        #del options
+
+        self.func = func
+        self.y0 = y0
+
+        if step_size is not None and grid_constructor is None:
+            self.grid_constructor = self._grid_constructor_from_step_size(step_size)
+        elif grid_constructor is None:
+            self.grid_constructor = lambda f, y0, t: t # Same time step as time interval
+        else:
+            raise ValueError("step_size and grid_constructor are exclusive arguments.")
+
+    def _grid_constructor_from_step_size(self, step_size):
+
+        def _grid_constructor(func, y0, t):
+            start_time = t[0]
+            end_time = t[-1]
+
+            niters = torch.ceil((end_time - start_time) / step_size + 1).item()
+            t_infer = torch.arange(0, niters).to(t) * step_size + start_time
+            if t_infer[-1] > t[-1]:
+                t_infer[-1] = t[-1]
+            return t_infer
+
+        return _grid_constructor
+
+    @property
+    @abc.abstractmethod
+    def order(self):
+        pass
+
+    @abc.abstractmethod
+    def step_func(self, func, t, dt, y):
+        pass
+
+    def integrate(self, t):
+        _assert_increasing(t)
+        t = t.type_as(self.y0[0]) # (n_time, )
+        time_grid = self.grid_constructor(self.func, self.y0, t)
+        #print('time_grid:', time_grid.size())
+        #print('t:', t.size())
+        assert time_grid[0] == t[0] and time_grid[-1] == t[-1]
+        time_grid = time_grid.to(self.y0[0])
+
+        solution = [self.y0]
+
+        j = 1
+        y0 = self.y0
+        for t0, t1 in zip(time_grid[:-1], time_grid[1:]):
+            dy = self.step_func(self.func, t0, t1 - t0, y0)
+            y1 = tuple(y0_ + dy_ for y0_, dy_ in zip(y0, dy))
+            y0 = y1
+            while j < len(t) and t1 >= t[j]:
+                solution.append(self._linear_interp(t0, t1, y0, y1, t[j]))
+                j += 1
+            '''if self.contours is not None:
+                if self.adjoint:
+                    for i in range(1, len(t)):
+                        ys = list(self._linear_interp(t0, t1, y0, y1, t[j])) # tuple: (y0, **back_grad) -> y0: (n_batch, spatial_shape)
+                        ys[0] = self.set_BC(ys[0], self.contours[:, i]) # (n_batch, 4 / 6, BC_size, c)
+                        solution.append(tuple(ys))
+                        j += 1
+                else:
+                    while j < len(t) and t1 >= t[j]:
+                        y = torch.stack(self._linear_interp(t0, t1, y0, y1, t[j])) # (n_batch, 1, spatial_shape)
+                        y = self.set_BC(y[:, 0], self.contours[:, j]).unsqueeze(1) 
+                        solution.append(tuple(y))
+                        j += 1
+            elif self.dcontours is not None:
+                if self.adjoint:
+                    for i in range(1, len(t)):
+                        ys = list(self._linear_interp(t0, t1, y0, y1, t[j])) # tuple: (y0, **back_grad) -> y0: (n_batch, spatial_shape)
+                        ys[0] = self.add_dBC(ys[0], self.dcontours[:, j]) # (n_batch, 4 / 6, BC_size, c)
+                        solution.append(tuple(ys))
+                else:
+                    while j < len(t) and t1 >= t[j]:
+                        y = torch.stack(self._linear_interp(t0, t1, y0, y1, t[j])) # (n_batch, 1, spatial_shape)
+                        y = self.add_dBC(y[:, 0], self.dcontours[:, j]).unsqueeze(1)
+                        solution.append(tuple(y))
+                        j += 1
+            else:
+                while j < len(t) and t1 >= t[j]:
+                    solution.append(self._linear_interp(t0, t1, y0, y1, t[j]))
+                    j += 1'''
+        return tuple(map(torch.stack, tuple(zip(*solution)))) # (batch, time)
+
+    def _linear_interp(self, t0, t1, y0, y1, t):
+        if t == t0:
+            return y0
+        if t == t1:
+            return y1
+        t0, t1, t = t0.to(y0[0]), t1.to(y0[0]), t.to(y0[0])
+        slope = tuple((y1_ - y0_) / (t1 - t0) for y0_, y1_, in zip(y0, y1))
+        return tuple(y0_ + slope_ * (t - t0) for y0_, slope_ in zip(y0, slope))

ShapeID/DiffEqs/tsit5.py

@@ -0,0 +1,139 @@
+import torch
+from ShapeID.DiffEqs.misc import _scaled_dot_product, _convert_to_tensor, _is_finite, _select_initial_step, _handle_unused_kwargs
+from ShapeID.DiffEqs.solvers import AdaptiveStepsizeODESolver
+from ShapeID.DiffEqs.rk_common import _RungeKuttaState, _ButcherTableau, _runge_kutta_step
+
+# Parameters from Tsitouras (2011).
+_TSITOURAS_TABLEAU = _ButcherTableau(
+    alpha=[0.161, 0.327, 0.9, 0.9800255409045097, 1., 1.],
+    beta=[
+        [0.161],
+        [-0.008480655492357, 0.3354806554923570],
+        [2.897153057105494, -6.359448489975075, 4.362295432869581],
+        [5.32586482843925895, -11.74888356406283, 7.495539342889836, -0.09249506636175525],
+        [5.86145544294642038, -12.92096931784711, 8.159367898576159, -0.071584973281401006, -0.02826905039406838],
+        [0.09646076681806523, 0.01, 0.4798896504144996, 1.379008574103742, -3.290069515436081, 2.324710524099774],
+    ],
+    c_sol=[0.09646076681806523, 0.01, 0.4798896504144996, 1.379008574103742, -3.290069515436081, 2.324710524099774, 0],
+    c_error=[
+        0.09646076681806523 - 0.001780011052226,
+        0.01 - 0.000816434459657,
+        0.4798896504144996 - -0.007880878010262,
+        1.379008574103742 - 0.144711007173263,
+        -3.290069515436081 - -0.582357165452555,
+        2.324710524099774 - 0.458082105929187,
+        -1 / 66,
+    ],
+)
+
+
+def _interp_coeff_tsit5(t0, dt, eval_t):
+    t = float((eval_t - t0) / dt)
+    b1 = -1.0530884977290216 * t * (t - 1.3299890189751412) * (t**2 - 1.4364028541716351 * t + 0.7139816917074209)
+    b2 = 0.1017 * t**2 * (t**2 - 2.1966568338249754 * t + 1.2949852507374631)
+    b3 = 2.490627285651252793 * t**2 * (t**2 - 2.38535645472061657 * t + 1.57803468208092486)
+    b4 = -16.54810288924490272 * (t - 1.21712927295533244) * (t - 0.61620406037800089) * t**2
+    b5 = 47.37952196281928122 * (t - 1.203071208372362603) * (t - 0.658047292653547382) * t**2
+    b6 = -34.87065786149660974 * (t - 1.2) * (t - 0.666666666666666667) * t**2
+    b7 = 2.5 * (t - 1) * (t - 0.6) * t**2
+    return [b1, b2, b3, b4, b5, b6, b7]
+
+
+def _interp_eval_tsit5(t0, t1, k, eval_t):
+    dt = t1 - t0
+    y0 = tuple(k_[0] for k_ in k)
+    interp_coeff = _interp_coeff_tsit5(t0, dt, eval_t)
+    y_t = tuple(y0_ + _scaled_dot_product(dt, interp_coeff, k_) for y0_, k_ in zip(y0, k))
+    return y_t
+
+
+def _optimal_step_size(last_step, mean_error_ratio, safety=0.9, ifactor=10.0, dfactor=0.2, order=5):
+    """Calculate the optimal size for the next Runge-Kutta step."""
+    if mean_error_ratio == 0:
+        return last_step * ifactor
+    if mean_error_ratio < 1:
+        dfactor = _convert_to_tensor(1, dtype=torch.float64, device=mean_error_ratio.device)
+    error_ratio = torch.sqrt(mean_error_ratio).type_as(last_step)
+    exponent = torch.tensor(1 / order).type_as(last_step)
+    factor = torch.max(1 / ifactor, torch.min(error_ratio**exponent / safety, 1 / dfactor))
+    return last_step / factor
+
+
+def _abs_square(x):
+    return torch.mul(x, x)
+
+
+class Tsit5Solver(AdaptiveStepsizeODESolver):
+
+    def __init__(
+        self, func, y0, rtol, atol, first_step=None, safety=0.9, ifactor=10.0, dfactor=0.2, max_num_steps=2**31 - 1,
+        **unused_kwargs
+    ):
+        _handle_unused_kwargs(self, unused_kwargs)
+        del unused_kwargs
+
+        self.func = func
+        self.y0 = y0
+        self.rtol = rtol
+        self.atol = atol
+        self.first_step = first_step
+        self.safety = _convert_to_tensor(safety, dtype=torch.float64, device=y0[0].device)
+        self.ifactor = _convert_to_tensor(ifactor, dtype=torch.float64, device=y0[0].device)
+        self.dfactor = _convert_to_tensor(dfactor, dtype=torch.float64, device=y0[0].device)
+        self.max_num_steps = _convert_to_tensor(max_num_steps, dtype=torch.int32, device=y0[0].device)
+
+    def before_integrate(self, t):
+        if self.first_step is None:
+            first_step = _select_initial_step(self.func, t[0], self.y0, 4, self.rtol, self.atol).to(t)
+        else:
+            first_step = _convert_to_tensor(0.01, dtype=t.dtype, device=t.device)
+        self.rk_state = _RungeKuttaState(
+            self.y0,
+            self.func(t[0].type_as(self.y0[0]), self.y0), t[0], t[0], first_step,
+            tuple(map(lambda x: [x] * 7, self.y0))
+        )
+
+    def advance(self, next_t):
+        """Interpolate through the next time point, integrating as necessary."""
+        n_steps = 0
+        while next_t > self.rk_state.t1:
+            assert n_steps < self.max_num_steps, 'max_num_steps exceeded ({}>={})'.format(n_steps, self.max_num_steps)
+            self.rk_state = self._adaptive_tsit5_step(self.rk_state)
+            n_steps += 1
+        return _interp_eval_tsit5(self.rk_state.t0, self.rk_state.t1, self.rk_state.interp_coeff, next_t)
+
+    def _adaptive_tsit5_step(self, rk_state):
+        """Take an adaptive Runge-Kutta step to integrate the DiffEqs."""
+        y0, f0, _, t0, dt, _ = rk_state
+        ########################################################
+        #                      Assertions                      #
+        ########################################################
+        assert t0 + dt > t0, 'underflow in dt {}'.format(dt.item())
+        for y0_ in y0:
+            assert _is_finite(torch.abs(y0_)), 'non-finite values in state `y`: {}'.format(y0_)
+        y1, f1, y1_error, k = _runge_kutta_step(self.func, y0, f0, t0, dt, tableau=_TSITOURAS_TABLEAU)
+
+        ########################################################
+        #                     Error Ratio                      #
+        ########################################################
+        error_tol = tuple(self.atol + self.rtol * torch.max(torch.abs(y0_), torch.abs(y1_)) for y0_, y1_ in zip(y0, y1))
+        tensor_error_ratio = tuple(y1_error_ / error_tol_ for y1_error_, error_tol_ in zip(y1_error, error_tol))
+        sq_error_ratio = tuple(
+            torch.mul(tensor_error_ratio_, tensor_error_ratio_) for tensor_error_ratio_ in tensor_error_ratio
+        )
+        mean_error_ratio = (
+            sum(torch.sum(sq_error_ratio_) for sq_error_ratio_ in sq_error_ratio) /
+            sum(sq_error_ratio_.numel() for sq_error_ratio_ in sq_error_ratio)
+        )
+        accept_step = mean_error_ratio <= 1
+
+        ########################################################
+        #                   Update RK State                    #
+        ########################################################
+        y_next = y1 if accept_step else y0
+        f_next = f1 if accept_step else f0
+        t_next = t0 + dt if accept_step else t0
+        dt_next = _optimal_step_size(dt, mean_error_ratio, self.safety, self.ifactor, self.dfactor)
+        k_next = k if accept_step else self.rk_state.interp_coeff
+        rk_state = _RungeKuttaState(y_next, f_next, t0, t_next, dt_next, k_next)
+        return rk_state

ShapeID/__init__.py

@@ -0,0 +1 @@
+from utils import *

ShapeID/demo2d.py

@@ -0,0 +1,102 @@
+# ported from https://github.com/pvigier/perlin-numpy
+
+import os, sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import time, datetime
+
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+
+from misc import stream_2D, V_plot
+from utils.misc import viewVolume, make_dir
+
+from perlin2d import *
+
+
+#from ShapeID.DiffEqs.odeint import odeint
+from ShapeID.DiffEqs.adjoint import odeint_adjoint as odeint
+from ShapeID.DiffEqs.pde import AdvDiffPDE
+
+
+ 
+
+if __name__ == '__main__':
+
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        
+    
+    image, mask_image = generate_perlin_noise_2d([256, 256], [2, 2], percentile = 80)  
+    plt.imshow(image, cmap='gray') #, interpolation='lanczos')
+    plt.axis('off')
+    plt.savefig('out/2d/image.png') 
+    plt.imshow(mask_image, cmap='gray') #, interpolation='lanczos')
+    plt.axis('off')
+    plt.savefig('out/2d/mask_image.png') 
+
+
+
+    curl, mask_curl = generate_perlin_noise_2d([256, 256], [2, 2], percentile = 80)
+    plt.imshow(curl, cmap='gray') #, interpolation='lanczos')
+    plt.axis('off')
+    plt.savefig('out/2d/curl.png') 
+    plt.imshow(mask_curl, cmap='gray') #, interpolation='lanczos')
+    plt.axis('off')
+    plt.savefig('out/2d/mask_curl.png') 
+     
+
+    dx, dy = stream_2D(torch.from_numpy(curl))
+    V_plot(dx.numpy(), dy.numpy(), 'out/2d/V.png')
+
+    plt.imshow(mask_image, cmap='gray') #, interpolation='lanczos')
+    plt.axis('off')
+    plt.savefig('out/2d/image_with_v.png')  
+    #plt.close()
+
+
+    dt = 0.15
+    nt = 21
+    integ_method = 'dopri5' # choices=['dopri5', 'adams', 'rk4', 'euler'] 
+    t = torch.from_numpy(np.arange(nt) * dt).to(device)
+    thres = 0.9
+
+    initial = torch.from_numpy(mask_image)
+    Vx, Vy = dx * 1000, dy * 1000
+
+    forward_pde = AdvDiffPDE(data_spacing=[1., 1.], 
+                             perf_pattern='adv', 
+                             V_type='vector_div_free', 
+                             V_dict={'Vx': Vx, 'Vy': Vy},
+                             BC='neumann', 
+                             dt=dt, 
+                             device=device
+                             )
+    
+  
+    start_time = time.time()
+    noise_progression = odeint(forward_pde, 
+                               initial.unsqueeze(0), 
+                               t, dt, method = integ_method
+                               )[:, 0]
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    print('Time {}'.format(total_time_str)) 
+    noise_progression = noise_progression[::2]
+
+
+    noise_progression = noise_progression.numpy()
+    make_dir('out/2d/progression')
+
+    for i, noise_t in enumerate(noise_progression): 
+        print(i, noise_t.mean())
+
+        noise_t[noise_t > thres] = 1
+        noise_t[noise_t <= thres] = 0
+
+        #fig = plt.figure()
+        plt.imshow(noise_t, cmap='gray') #, interpolation='lanczos')
+        plt.savefig('out/2d/progression/%d.png' % i) 
+        #plt.close() 
+
+    viewVolume(noise_progression, names = ['noise_progression'], save_dir = 'out/2d/progression')

ShapeID/demo3d.py

@@ -0,0 +1,91 @@
+# ported from https://github.com/pvigier/perlin-numpy
+
+import os, sys
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import time, datetime
+
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+
+from misc import stream_3D, V_plot, center_crop
+from utils.misc import viewVolume, make_dir, read_image
+
+
+#from ShapeID.DiffEqs.odeint import odeint
+from ShapeID.DiffEqs.adjoint import odeint_adjoint as odeint
+from ShapeID.DiffEqs.pde import AdvDiffPDE
+
+from perlin3d import *
+
+
+
+
+if __name__ == '__main__': 
+
+    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+
+    percentile = 80
+    
+
+    #image, mask_image = generate_perlin_noise_3d([128, 128, 128], [2, 2, 2], tileable=(True, False, False), percentile = percentile)
+    #viewVolume(image, names = ['image'], save_dir = '/autofs/space/yogurt_003/users/pl629/code/MTBrainID/ShapeID/out/3d')
+    #viewVolume(mask_image, names = ['mask_image'], save_dir = '/autofs/space/yogurt_003/users/pl629/code/MTBrainID/ShapeID/out/3d') 
+
+
+    #mask_image, aff = read_image('/autofs/space/yogurt_001/users/pl629/data/adni/pathology_probability/subject_193441.nii.gz')
+    mask_image, aff = read_image('/autofs/space/yogurt_001/users/pl629/data/isles2022/pathology_probability/sub-strokecase0127.nii.gz')
+    mask_image, _, _ = center_crop(torch.from_numpy(mask_image), win_size = [128, 128, 128])
+    mask_image = mask_image[0, 0].numpy()
+
+    shape = mask_image.shape
+
+    curl_a, _ = generate_perlin_noise_3d(shape, [2, 2, 2], tileable=(True, False, False), percentile = percentile) 
+    curl_b, _ = generate_perlin_noise_3d(shape, [2, 2, 2], tileable=(True, False, False), percentile = percentile) 
+    curl_c, _ = generate_perlin_noise_3d(shape, [2, 2, 2], tileable=(True, False, False), percentile = percentile) 
+    dx, dy, dz = stream_3D(torch.from_numpy(curl_a), torch.from_numpy(curl_b), torch.from_numpy(curl_c)) 
+
+
+    dt = 0.1
+    nt = 10
+    integ_method = 'dopri5' # choices=['dopri5', 'adams', 'rk4', 'euler'] 
+    t = torch.from_numpy(np.arange(nt) * dt).to(device)
+    thres = 0.5
+
+    initial = torch.from_numpy(mask_image)[None] # (batch=1, h, w)
+    Vx, Vy, Vz = dx * 500, dy * 500, dz * 500
+    print(abs(Vx).mean(), abs(Vy).mean(), abs(Vz).mean())
+
+    forward_pde = AdvDiffPDE(data_spacing=[1., 1., 1.], 
+                             perf_pattern='adv', 
+                             V_type='vector_div_free', 
+                             V_dict={'Vx': Vx, 'Vy': Vy, 'Vz': Vz},
+                             BC='neumann', 
+                             dt=dt, 
+                             device=device
+                             )
+
+    
+    start_time = time.time()
+    noise_progression = odeint(forward_pde, 
+                               initial, 
+                               t, dt, method = integ_method
+                               )[:, 0] # (nt, n_batch, h, w)
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    print('Time {}'.format(total_time_str)) 
+    
+    noise_progression = noise_progression[::2]
+    noise_progression = noise_progression.numpy()
+    make_dir('out/3d/progression')
+
+
+    for i, noise_t in enumerate(noise_progression):
+        noise_t[noise_t > 1] = 1
+        noise_t[noise_t <= thres] = 0
+        print(i, noise_t.mean())
+        viewVolume(noise_t, names = ['noise_%s' % i], save_dir = '/autofs/space/yogurt_003/users/pl629/code/MTBrainID/ShapeID/out/3d/progression')
+        
+        noise_t[noise_t > 0.] = 1
+        viewVolume(noise_t, names = ['noise_%s_mask' % i], save_dir = '/autofs/space/yogurt_003/users/pl629/code/MTBrainID/ShapeID/out/3d/progression')

ShapeID/misc.py

@@ -0,0 +1,261 @@
+# ported from https://github.com/pvigier/perlin-numpy
+
+
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+
+def center_crop(img, win_size = [220, 220, 220]):
+    # center crop
+    if len(img.shape) == 4: 
+        img = torch.permute(img, (3, 0, 1, 2)) # (move last dim to first)
+        img = img[None]
+        permuted = True
+    else: 
+        assert len(img.shape) == 3
+        img = img[None, None]
+        permuted = False
+
+    orig_shp = img.shape[2:] # (1, d, s, r, c)
+    if win_size is None:
+        if permuted:
+            return torch.permute(img, (0, 2, 3, 4, 1)), [0, 0, 0], orig_shp
+        return img, [0, 0, 0], orig_shp
+    elif orig_shp[0] > win_size[0] or orig_shp[1] > win_size[1] or orig_shp[2] > win_size[2]:
+        crop_start = [ max((orig_shp[i] - win_size[i]), 0) // 2 for i in range(3) ]
+        crop_img = img[ :, :, crop_start[0] : crop_start[0] + win_size[0], 
+                   crop_start[1] : crop_start[1] + win_size[1], 
+                   crop_start[2] : crop_start[2] + win_size[2]]
+        if permuted:
+            return torch.permute(crop_img, (0, 2, 3, 4, 1)), [0, 0, 0], orig_shp
+        return crop_img, crop_start, orig_shp
+    else:
+        if permuted:
+            return torch.permute(img, (0, 2, 3, 4, 1)), [0, 0, 0], orig_shp
+        return img, [0, 0, 0], orig_shp
+
+
+
+def V_plot(Vx, Vy, save_path):
+    # Meshgrid 
+    X,Y = np.meshgrid(np.arange(0, Vx.shape[0], 1), np.arange(0, Vx.shape[1], 1)) 
+    # Assign vector directions 
+    Ex = Vx 
+    Ey = Vy 
+
+    # Depict illustration 
+    plt.figure() 
+    plt.streamplot(X,Y,Ex,Ey, density=1.4, linewidth=None, color='orange')  
+    plt.axis('off')
+    plt.savefig(save_path)
+    #plt.show()
+
+def stream_2D(Phi, batched = False, delta_lst = [1., 1.]):
+    '''
+    input: Phi as a scalar field in 2D grid: (r, c) or (n_batch, r, c)
+    output: curl of Phi (divergence-free by definition)
+    ''' 
+    dD = gradient_c(Phi, batched = batched, delta_lst = delta_lst) 
+    Vx = - dD[..., 1]
+    Vy = dD[..., 0]
+    return Vx, Vy
+
+
+def stream_3D(Phi_a, Phi_b, Phi_c, batched = False, delta_lst = [1., 1., 1.]):
+    '''
+    input: (batch, s, r, c)
+    '''
+    device = Phi_a.device
+    dDa = gradient_c(Phi_a, batched = batched, delta_lst = delta_lst)
+    dDb = gradient_c(Phi_b, batched = batched, delta_lst = delta_lst)
+    dDc = gradient_c(Phi_c, batched = batched, delta_lst = delta_lst)
+    Va_x, Va_y, Va_z = dDa[..., 0], dDa[..., 1], dDa[..., 2]
+    Vb_x, Vb_y, Vb_z = dDb[..., 0], dDb[..., 1], dDb[..., 2]
+    Vc_x, Vc_y, Vc_z = dDc[..., 0], dDc[..., 1], dDc[..., 2]
+    Vx = Vc_y - Vb_z
+    Vy = Va_z - Vc_x
+    Vz = Vb_x - Va_y
+    return Vx, Vy, Vz
+
+
+
+def gradient_f(X, batched = False, delta_lst = [1., 1., 1.]):
+    '''
+    Compute gradient of a torch tensor "X" in each direction
+    Upper-boundaries: Backward Difference
+    Non-boundaries & Upper-boundaries: Forward Difference
+    if X is batched: (n_batch, ...);
+    else: (...)
+    '''
+    device = X.device
+    dim = len(X.size()) - 1 if batched else len(X.size())
+    #print(batched)
+    #print(dim)
+    if dim == 1:
+        #print('dim = 1')
+        dX = torch.zeros(X.size(), dtype = torch.float, device = device)
+        X = X.permute(1, 0) if batched else X
+        dX = dX.permute(1, 0) if batched else dX
+        dX[-1] = X[-1] - X[-2] # Backward Difference
+        dX[:-1] = X[1:] - X[:-1] # Forward Difference
+
+        dX = dX.permute(1, 0) if batched else dX
+        dX /= delta_lst[0]
+    elif dim == 2:
+        #print('dim = 2')
+        dX = torch.zeros(X.size() + tuple([2]), dtype = torch.float, device = device)
+        X = X.permute(1, 2, 0) if batched else X
+        dX = dX.permute(1, 2, 3, 0) if batched else dX # put batch to last dim
+        dX[-1, :, 0] = X[-1, :] - X[-2, :] # Backward Difference
+        dX[:-1, :, 0] = X[1:] - X[:-1] # Forward Difference
+
+        dX[:, -1, 1] = X[:, -1] - X[:, -2] # Backward Difference
+        dX[:, :-1, 1] = X[:, 1:] - X[:, :-1] # Forward Difference
+
+        dX = dX.permute(3, 0, 1, 2) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+    elif dim == 3:
+        #print('dim = 3')
+        dX = torch.zeros(X.size() + tuple([3]), dtype = torch.float, device = device)
+        X = X.permute(1, 2, 3, 0) if batched else X
+        dX = dX.permute(1, 2, 3, 4, 0) if batched else dX
+        dX[-1, :, :, 0] = X[-1, :, :] - X[-2, :, :] # Backward Difference
+        dX[:-1, :, :, 0] = X[1:] - X[:-1] # Forward Difference
+
+        dX[:, -1, :, 1] = X[:, -1] - X[:, -2] # Backward Difference
+        dX[:, :-1, :, 1] = X[:, 1:] - X[:, :-1] # Forward Difference
+
+        dX[:, :, -1, 2] = X[:, :, -1] - X[:, :, -2] # Backward Difference
+        dX[:, :, :-1, 2] = X[:, :, 1:] - X[:, :, :-1] # Forward Difference
+
+        dX = dX.permute(4, 0, 1, 2, 3) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+        dX[..., 2] /= delta_lst[2]
+    return dX
+
+
+def gradient_b(X, batched = False, delta_lst = [1., 1., 1.]):
+    '''
+    Compute gradient of a torch tensor "X" in each direction
+    Non-boundaries & Upper-boundaries: Backward Difference
+    Lower-boundaries: Forward Difference
+    if X is batched: (n_batch, ...);
+    else: (...)
+    '''
+    device = X.device
+    dim = len(X.size()) - 1 if batched else len(X.size())
+    #print(batched)
+    #print(dim)
+    if dim == 1:
+        #print('dim = 1')
+        dX = torch.zeros(X.size(), dtype = torch.float, device = device)
+        X = X.permute(1, 0) if batched else X
+        dX = dX.permute(1, 0) if batched else dX
+        dX[1:] = X[1:] - X[:-1] # Backward Difference
+        dX[0] = X[1] - X[0] # Forward Difference
+
+        dX = dX.permute(1, 0) if batched else dX
+        dX /= delta_lst[0]
+    elif dim == 2:
+        #print('dim = 2')
+        dX = torch.zeros(X.size() + tuple([2]), dtype = torch.float, device = device)
+        X = X.permute(1, 2, 0) if batched else X
+        dX = dX.permute(1, 2, 3, 0) if batched else dX # put batch to last dim
+        dX[1:, :, 0] = X[1:, :] - X[:-1, :] # Backward Difference
+        dX[0, :, 0] = X[1] - X[0] # Forward Difference
+
+        dX[:, 1:, 1] = X[:, 1:] - X[:, :-1] # Backward Difference
+        dX[:, 0, 1] = X[:, 1] - X[:, 0] # Forward Difference
+
+        dX = dX.permute(3, 0, 1, 2) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+    elif dim == 3:
+        #print('dim = 3')
+        dX = torch.zeros(X.size() + tuple([3]), dtype = torch.float, device = device)
+        X = X.permute(1, 2, 3, 0) if batched else X
+        dX = dX.permute(1, 2, 3, 4, 0) if batched else dX
+        dX[1:, :, :, 0] = X[1:, :, :] - X[:-1, :, :] # Backward Difference
+        dX[0, :, :, 0] = X[1] - X[0] # Forward Difference
+
+        dX[:, 1:, :, 1] = X[:, 1:] - X[:, :-1] # Backward Difference
+        dX[:, 0, :, 1] = X[:, 1] - X[:, 0] # Forward Difference
+
+        dX[:, :, 1:, 2] = X[:, :, 1:] - X[:, :, :-1] # Backward Difference
+        dX[:, :, 0, 2] = X[:, :, 1] - X[:, :, 0] # Forward Difference
+
+        dX = dX.permute(4, 0, 1, 2, 3) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+        dX[..., 2] /= delta_lst[2]
+    return dX
+  
+
+def gradient_c(X, batched = False, delta_lst = [1., 1., 1.]):
+    '''
+    Compute gradient of a torch tensor "X" in each direction
+    Non-boundaries: Central Difference
+    Upper-boundaries: Backward Difference
+    Lower-boundaries: Forward Difference
+    if X is batched: (n_batch, ...);
+    else: (...)
+    ''' 
+
+    device = X.device
+    dim = len(X.size()) - 1 if batched else len(X.size())
+    #print(X.size())
+    #print(batched)
+    #print(dim)
+    if dim == 1:
+        #print('dim = 1')
+        dX = torch.zeros(X.size(), dtype = torch.float, device = device)
+        X = X.permute(1, 0) if batched else X
+        dX = dX.permute(1, 0) if batched else dX
+        dX[1:-1] = (X[2:] - X[:-2]) / 2 # Central Difference
+        dX[0] = X[1] - X[0] # Forward Difference
+        dX[-1] = X[-1] - X[-2] # Backward Difference
+
+        dX = dX.permute(1, 0) if batched else dX
+        dX /= delta_lst[0]
+    elif dim == 2:
+        #print('dim = 2')
+        dX = torch.zeros(X.size() + tuple([2]), dtype = torch.float, device = device)
+        X = X.permute(1, 2, 0) if batched else X
+        dX = dX.permute(1, 2, 3, 0) if batched else dX # put batch to last dim
+        dX[1:-1, :, 0] = (X[2:, :] - X[:-2, :]) / 2
+        dX[0, :, 0] = X[1] - X[0]
+        dX[-1, :, 0] = X[-1] - X[-2]
+        dX[:, 1:-1, 1] = (X[:, 2:] - X[:, :-2]) / 2
+        dX[:, 0, 1] = X[:, 1] - X[:, 0]
+        dX[:, -1, 1] = X[:, -1] - X[:, -2]
+
+        dX = dX.permute(3, 0, 1, 2) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+    elif dim == 3:
+        #print('dim = 3')
+        dX = torch.zeros(X.size() + tuple([3]), dtype = torch.float, device = device)
+        X = X.permute(1, 2, 3, 0) if batched else X
+        dX = dX.permute(1, 2, 3, 4, 0) if batched else dX
+        dX[1:-1, :, :, 0] = (X[2:, :, :] - X[:-2, :, :]) / 2
+        dX[0, :, :, 0] = X[1] - X[0]
+        dX[-1, :, :, 0] = X[-1] - X[-2]
+        dX[:, 1:-1, :, 1] = (X[:, 2:, :] - X[:, :-2, :]) / 2
+        dX[:, 0, :, 1] = X[:, 1] - X[:, 0]
+        dX[:, -1, :, 1] = X[:, -1] - X[:, -2]
+        dX[:, :, 1:-1, 2] = (X[:, :, 2:] - X[:, :, :-2]) / 2
+        dX[:, :, 0, 2] = X[:, :, 1] - X[:, :, 0]
+        dX[:, :, -1, 2] = X[:, :, -1] - X[:, :, -2]
+
+        dX = dX.permute(4, 0, 1, 2, 3) if batched else dX
+        dX[..., 0] /= delta_lst[0]
+        dX[..., 1] /= delta_lst[1]
+        dX[..., 2] /= delta_lst[2]
+     
+    return dX
+
+