print('Regularization: ' + str(ADVERSARIAL_REGULARIZATION))
with mrcfile.open(
file['external_reconstruct_general']['rlnExtReconsDataReal']) as mrc:
data_real = mrc.data
with mrcfile.open(
file['external_reconstruct_general']['rlnExtReconsDataImag']) as mrc:
data_im = mrc.data
with mrcfile.open(
file['external_reconstruct_general']['rlnExtReconsWeight']) as mrc:
kernel = mrc.data
target_path = file['external_reconstruct_general']['rlnExtReconsResult']
complex_data = data_real + 1j * data_im
regularizer = AdversarialRegulariser(SAVES_PATH)
tikhonov_kernel = kernel + REGULARIZATION_TY
precondioner = np.abs(np.divide(1, tikhonov_kernel))
precondioner /= precondioner.max()
tikhonov = np.divide(complex_data, tikhonov_kernel)
reco = np.copy(tikhonov)
# The scales produce gradients of order 1
ADVERSARIAL_SCALE = (96**(-0.5))
DATA_SCALE = 1 / (10 * 96**3)
IMAGING_SCALE = 96
for k in range(70):
STEP_SIZE = 1.0 * 1 / np.sqrt(1 + k / 20)
GAMMA = 5.0
CUTOFF = 20.0
def data_augmentation(gt, adv):
_, adv1 = interpolation(gt, adv)
_, adv2 = phase_augmentation(gt, adv1)
_, adv3 = positivity(gt, adv2)
new_gt, new_adv = rotation_translation(gt, adv3)
return new_gt, new_adv
saves_path = '/local/scratch/public/sl767/SPA/Saves/Adversarial_Regulariser/Cutoff_20/Translation_Augmentation'
regularizer = AdversarialRegulariser(saves_path,
data_augmentation,
s=S,
cutoff=20.0,
lmb=LMB,
gamma=GAMMA)
def evaluate():
gt, adv = get_batch(eval_data=True,
noise_levels=['01', '016'],
methods=['EM', 'SGD'])
regularizer.test(groundTruth=gt, adversarial=adv)
def train(steps):
for k in range(steps):
gt, adv = get_batch(eval_data=False)
regularizer.train(groundTruth=gt,
y = tf.spectral.rfft3d(gt[..., 0])
phase = 2 * np.pi * tf.random_uniform(
shape=tf.shape(y), minval=0, maxval=1)
com_phase = tf.exp(1j * tf.cast(phase, tf.complex64))
y = tf.multiply(com_phase, y)
adv_phase = tf.expand_dims(tf.spectral.irfft3d(y), axis=-1)
eps1 = tf.random_uniform(shape=(BATCH_SIZE, 1, 1, 1, 1),
minval=0,
maxval=0.5)
adv_new = tf.multiply(eps1, adv_phase) + tf.multiply(
tf.ones(shape=(BATCH_SIZE, 1, 1, 1, 1)) - eps1, adv_inter)
return gt, adv_new
saves_path = '/local/scratch/public/sl767/SPA/Saves/Adversarial_Regulariser/SGD_Trained/phase_augmentation/'
regularizer = AdversarialRegulariser(saves_path, data_augmentation)
def evaluate():
gt, adv = get_batch(training_data=True)
regularizer.test(groundTruth=gt, adversarial=adv, fourier_data=False)
def train(steps):
for k in range(steps):
gt, adv = get_batch()
regularizer.train(groundTruth=gt,
adversarial=adv,
learning_rate=LEARNING_RATE,
fourier_data=False)
if k % 10 == 0:
def data_augmentation(gt, adv):#, noise_lvl):
# _, adv = interpolation(gt, adv)
# _, adv = phase_augmentation(gt, adv)
if TRAIN_ON_POS:
# print('hello')
_, adv = positivity(gt, adv)
### Do 90-rots in numpy instead
# new_gt, new_adv = rotation_translation(gt, adv)
# new_gt, new_adv = new_gt/noise_lvl, new_adv/noise_lvl
# new_gt, new_adv = new_gt/500, new_adv/500 # Bring back to old scale for training
return gt, adv
regularizer = AdversarialRegulariser(SAVES_PATH, data_augmentation,
s=S, cutoff=20.0, lmb=LMB, gamma=GAMMA)
#log_file = open(LOG_FILE_PATH, "w")
#log_file.write('Train data:\n' + str(TRAIN_DICT) + '\n')
#log_file.write('Eval data:\n' + str(EVAL_DICT) + '\n')
#log_file.close()
def evaluate():
gt, adv, nl = get_batch(batch_size=BATCH_SIZE, noise_levels=EVAL_NOISE_LEVELS,
methods=EVAL_METHODS, data_dict=EVAL_DICT,
eval_data=True)
# noise_lvl = int(nl) / 100
regularizer.test(groundTruth=gt, adversarial=adv)#, noise_lvl=noise_lvl)
from ClassFiles.relion_fixed_it import load_star
from ClassFiles.AdversarialRegularizer import AdversarialRegulariser
from ClassFiles.ut import l2
from ClassFiles.Utilities import registration
from ClassFiles.ut import locate_gt, rfft, irfft, unify_form, Rescaler
# In[3]:
# saves_path = '/local/scratch/public/sl767/SPA/Saves/Adversarial_Regulariser/AllData/AllAugmentation'
PLATFORM_NODE = platform.node()
if PLATFORM_NODE == 'gelfand':
saves_path = '/mnt/data/zickert/'
elif PLATFORM_NODE == 'motel':
saves_path = '/local/scratch/public/sl767/'
saves_path += 'SPA/Saves/Adversarial_Regulariser/Cutoff_20/Translation_Augmentation'
regularizer = AdversarialRegulariser(saves_path)
# In[5]:
NUM_VOX = 150
TARGET_NUM_VOX = 96
IMAGING_SCALE = 300
def vis(data, fourier=True):
if fourier:
data = irfft(data, scaling=NUM_VOX**2)
slice_n = int(data.shape[0] // 2)
plt.imshow(IMAGING_SCALE * data.squeeze()[..., slice_n])
BASE_PATH = '/mnt/datahd/zickert/'
REGULARIZATION_TIK = 1e-3 # For initialization
PRECOND = False
REPORT = 10
SAVE_RECO_TO_DISK = False
EVAL_METRIC = 'L2_and_SSIM' # 'masked_L2' # 'masked_FSC'
NOISE_LEVEL = '01'
WIN_SIZE = 7 # for SSIM
REGULARIZATION_TIK_GD = 1e6 # During gradient descent
#%%
AR_PATH = '/mnt/datahd/zickert/SPA/Saves/SimDataPaper/Adversarial_Regulariser/no_fancy_data_aug_lr_5e-05_s_{}_pos_{}'.format(
SOBOLEV_CST, POS_AR)
regularizer = AdversarialRegulariser(AR_PATH)
POS_AR = ('1' == POS_AR)
print('Global parameters:')
print(SOBOLEV_CST, REGULARIZATION_TIK, PRECOND, REPORT, POS_AR, EVAL_METRIC,
NOISE_LEVEL)
print('################################')
#%%
#def runCommand(cmd_string, file_path=None):
# if file_path is None:
# sp.call(cmd_string.split(' '))
# else:
# file = open(file_path, 'w')
# sp.call(cmd_string.split(' '), stdout=file)
# file.close()