fmo_optim.py

import scipy
import logging
import argparse
import scipy.io as sio
from scipy.sparse import load_npz, \
				csc_matrix, save_npz, vstack
from os.path import join, expanduser
from .utils.optim import FMO

logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.DEBUG)
LOGGER = logging.getLogger(__name__)


parser = argparse.ArgumentParser(description='dose_parser')
parser.add_argument('--case', '-ca', type=str, default='case031',
						help='patient case' )
args = parser.parse_args()


# path where all the dose matrices are located: This is on the T drive
dosage_folder = join(expanduser('~'), 'mount', 'tdrive', \
					'Phys Research', 'Users', 'Lekan', 'radoncol', 'optimization')
os.chdir(dosage_folder)


"""
Prostate cases are from case031 throug case059 in ascending order
"""
NANGLE = 72

def load_mats(args, case_name, mask_indices=None):
	"""
	casename is the patient of interest
	mask indices is an optional variable (for 2D slice indices)
	"""

	global NANGLE # number of angles in total
	LOGGER.debug(' Loading patient {}'.format(case_name))


	# casenameos.chdir(dosage_folder, casename)
	nlp 	= sio.loadmat(case_name + '/nlpdata_' +str(NANGLE) + '.mat')
	params 	= sio.loadmat(case_name + '/parameters.mat')
	rect 	= sio.loadmat(case_name + '/' + case_name  + '_6MV_rect.mat')

	voxe       = nlp['voxe']#.astype(np.int64)
	bixe       = nlp['bixe']#.astype(np.int64)
	dijs       = nlp['dijs']  # this is float64 already
	numvoxels  = nlp['numvoxels'][0][0]+1 # cause python is 0 -indexed
	nbixel     = nlp['nbixel'][0][0]+1

	dMat = scipy.sparse.csc_matrix((dijs.squeeze(), (voxe.squeeze(), \
				bixe.squeeze())), shape=(numvoxels, nbixel), \
				dtype=np.float64, copy=False)

	"""
	This below optionally loads the dMat that you calculated in Matlab
	uncomment if you do not want to use the one that we calculated from scipy
	"""
	# dMat = sio.loadmat(case_name + '/'  + case_name \
	# 					+ '_dMat.mat')['dMat']

	# LOGGER.info('dMat shape: {}'.format(dMat.shape))

	# everything below creates the doses for 2D slices if we are optimizing for
	# 2D slices. You can ignore them for now
	if mask_indices:
		mask   	  =  rect['themask']
		pickstructs = rect['pickstructs'][0]

		def stack_arrays(arr1, arr2):
			return vstack((arr1, arr2))

		for slyce in mask_indices:
			structnnz, slice_dose=[], []
			IDXStructslice = []
			mask_ = mask.transpose()
			for i in range(len(pickstructs)):
				idxmask = np.zeros(mask_.shape)
				idxmask[np.nonzero(mask_)] = range(np.count_nonzero(mask_))
				idxmask=idxmask.transpose()
				idxslice = idxmask[:,:,slyce]

				idxstructslice = idxslice[np.nonzero(
									np.bitwise_and(mask[:,:, slyce], 2**i)
									)]
				IDXStructslice.append(idxstructslice)
				organ_dose = dMat[idxstructslice,:]

				structnnz.append(organ_dose.shape[0])
				slice_dose.append(organ_dose)
				if i > 0:
					slice_dose[i] = stack_arrays(slice_dose[i-1], slice_dose[i])
					# slice_dose[i] = cat_csc_sparse(slice_dose[i-1], slice_dose[i])
					IDXStructslice[i] = np.r_[IDXStructslice[i-1], IDXStructslice[i]]

				assert sum(structnnz) == slice_dose[-1].shape[0], 'nonzero indices are incorrect for slice and structnnz shapes'
			LOGGER.debug(' saving dose matrix for slice %d in %s', slyce, case_name)

			filename1 = join(dosage_folder, case_name, 'dose_slice_' + str(slyce)+ '.npz')
			if os.path.isfile(filename1):
				os.remove(filename1)
			save_npz(join(dosage_folder, case_name) + '/dose_slice_' + str(slyce)+ \
						'.npz', slice_dose[-1])
			if os.path.isfile(join(dosage_folder, case_name) + '/dose_slice_' + str(slyce)+ \
						'.mat'):
				os.remove(join(dosage_folder, case_name) + '/dose_slice_' + str(slyce)+ \
							'.mat')
			sio.savemat(join(dosage_folder, case_name) + '/dose_slice_' + str(slyce)+ \
						'.mat', {'dose_slice_'+ str(slyce): slice_dose[-1]}, oned_as='column')
			if os.path.isfile(join(dosage_folder, case_name) + '/slice_struct_' + str(slyce)+ \
						'.npz'):
				os.remove(join(dosage_folder, case_name) + '/slice_struct_' + str(slyce)+ \
							'.npz')
			np.savez(join(dosage_folder, case_name) + '/slice_struct_' + str(slyce)+ \
						'.npz', slice_struct=IDXStructslice[-1], order='F')


	bundle = dict(	nlp		=	nlp,
					params	=	params,
					rect    = 	rect,
					dMat	=	dMat)
	return bundle

def find_costs(args, dose_slice, bmindex, thresh,
				impupp, implow, structs_whole,
				 angles, slice_struct=None,
				 caseIdx='case031', sliceIdx=38):
	"""
		Everythiong here was translated from your projects.m file

		[dose_slice]: dMat or 2D dMat

		From your dataset:
		[bmindex]: obtained from main
		[thresh]: obtained from main
		[impupp]: loaded in main
		[structs_whole]: structs for the 3D patient
		[angles]: angles for which we are optimizing the FMO, arbitrary for testing code
		[slice_struct]: struct that corresponds to a 2D slice. This was precalculated in load_mats
		and loaded into at run time if we are considering 2D masks

		[caseIdx]: index of the patient we are considering e.g. case059
		sliceidx: Slice of patient within caseIdx


	"""
	global NANGLE

	def flatten(mat):
		if mat.ndim > 1:
			mat =  mat.flatten()
		return mat

	numvoxels    = dose_slice.shape[0]
	nbixel 		 = dose_slice.shape[-1]
	alphaPlusVec = np.zeros([numvoxels,1])
	alphaMinuVec = np.zeros([numvoxels,1])
	mVec         = np.zeros([numvoxels,1])

	if slice_struct is not None:
		slice_struct = slice_struct.astype(np.int64)
		if np.isnan(np.any(slice_struct)):
			LOGGER.critical('NANs encountered in slice_struct index. Please check this in find_costs')
		structs = structs_whole[slice_struct]

		assert structs.shape[0] == dose_slice.shape[0], 'structs and dose_slice need to be same shape along first dimension for 2D slice'

		if np.isnan(np.any(structs)):
			LOGGER.critical('NANs encountered in structs index. Please check this in find_costs')
	else:
		structs = structs_whole # for batch of patients

	for istructs in range(1,33):
		alphaPlusVec[np.where(structs == istructs)] = impupp[:,istructs]
		alphaMinuVec[np.where(structs == istructs)] = implow[:,istructs]
		mVec[np.where(structs == istructs),:] = thresh[:,istructs]

	alphaPlusVec = alphaPlusVec>0
	alphaMinuVec = alphaMinuVec>0

	nbixPerAngle = np.zeros([NANGLE,1])
	for iangle in range(NANGLE):
		nbixPerAngle[iangle, :] = np.nonzero(bmindex[:,:,iangle])[0].shape[0]
		# nbixPerAngle[iangle, :] = np.amax(
		# 							bmindex[:,:,iangle][np.nonzero(
		# 							bmindex[:,:,iangle])].shape)
	cumNbixPerAngle = np.append([0], np.cumsum(nbixPerAngle)).astype(np.int64)

	maskVec = np.zeros([nbixel,1])

	for i in range(len(angles)):
		idx2 = cumNbixPerAngle[angles[i]] # account for python/matlab indexing
		idx3 = cumNbixPerAngle[angles[i]+1]
		idxdiff = idx3 - cumNbixPerAngle[angles[i]]

		maskVec[range(idx2, idx3)] =  np.ones([idxdiff, 1])


	xVec 		= np.zeros([nbixel,1])
	Mdvh, Ndvh 	= DVHC.shape
	dVec 		= np.zeros([Mdvh,1])

	"""
	 the # commented codes below loads the calculated maskVec, mVec, alpha vectors
	 that you calculated in matlab and it has been verified that feeding these arrays
	 to the optimization engine works. The variables calculatred above in python is
	 where I am having issues with numpy and python 0-based indexing
	"""
	# maskVec = sio.loadmat(caseIdx + '/' + 'maskVec.mat')['maskVec']
	# mVec = sio.loadmat(caseIdx + '/' + 'mVec.mat')['mVec']
	# alphaPlusVec = sio.loadmat(caseIdx + '/' + 'alphaPlusVec.mat')['alphaPlusVec']
	# alphaMinuVec = sio.loadmat(caseIdx + '/' + 'alphaMinuVec.mat')['alphaMinuVec']

	fmo = FMO(args)  # this runs runGDM algorithm
	objValue, dose, xVec = fmo.optim(args, dose_slice, mVec, alphaPlusVec, \
					alphaMinuVec, xVec, maskVec, structs, TOL,\
					LAMBDA, DVHC, dVec, inistep=0, ifbreg=0)

	return objValue, dose, xVec


def main(args, caseIdx):

	bundle =  load_mats(args, caseIdx)
	bmindex, params, dMat		= bundle['nlp']['bmindex'], \
								 bundle['params'], bundle['dMat']
	mask   						= bundle['rect']['themask']

	structs    = params['structs'] if 'structs' in params.keys() \
								else bundle['nlp']['structs']
	thresh, impupp, implow     = params['thresh'], params['impupp'], \
								 params['implow']

	# fictitious angles for Testing
	angles = np.arange(5, 30, 5)
	angles = np.expand_dims(angles, axis=1)

	# find the indices that correspond the chosen angles [will be from 1 - 72]
	angle_indices = list(map(lambda x: np.abs((x/5)).astype(np.int64), \
								angles))

	cost, optimized_dose, xVec = find_costs(args, dMat, bmindex, \
									thresh, impupp, implow,  \
									structs, angle_indices, slice_struct=None,
									caseIdx=caseIdx, sliceIdx=slyce)

if __name__ == '__main__':
	chosen_case = 'case031' or args.case # as an example | prostate
	main(args, chosen_case)