Applying differentiable affine transformation


I have 4 anatomical views, one is 3D (sa) and the other ones are 2d(_la). I want to perform the following operation(multiview_affine) in a differentiable way. The problems that I am facing are mainly that I have to deal with a non integer operations and that a lot of points between 3D and 2D views lay outside. So I solve that iterating over all the 3d view points, computing the position of that points on each one of the other views and retriving its value, verifying that the float positions are converted to integers. I am not sure if this is the optimal way to do that, an dof course I don’t know how to write that on a pytorch differentiable neural network.

from nibabel.affines import apply_affine
import nibabel as nib
import numpy as np

def apply_affine_sa_la(_sa,_la,idx):

    sa2la_affine = np.linalg.inv(_la.affine).dot(_sa.affine)

    sape = _sa.get_fdata().shape

    data_la = _la .get_fdata()[...,idx]

    ph = np.zeros_like( _sa.get_fdata()[...,idx])

    for x in range(sape[0]):

        for y in range(sape[1]):

            for z in range(sape[2]):

                la_pos = apply_affine(sa2la_affine,np.array([x,y,z]))

                #la_pos = np.floor(la_pos)

                if np.abs(la_pos[2]) < 0.5:


                        ph[x,y,z] = data_la[int(la_pos[0]),int(la_pos[1]),0]





    return ph

def multiview_affine(_sa,_4c,_3c,_2c,idx):

    ph_0 = apply_affine_sa_la(_sa,_4c,idx)

    ph_1 = apply_affine_sa_la(_sa,_3c,idx)

    ph_2 = apply_affine_sa_la(_sa,_2c,idx)

    d = np.stack([ph_0,ph_1,ph_2],axis = 3)

    ph = d.max(axis=3)

    return ph