utils_sortem.py 17.8 KB
Newer Older
1

Jonathan Juhl's avatar
Jonathan Juhl committed
2
3


4
import tensorflow as tf
Jonathan Juhl's avatar
Jonathan Juhl committed
5
6
import matplotlib
matplotlib.use('Agg')
7
8
import numpy as np
from os.path import join
Jonathan Juhl's avatar
Jonathan Juhl committed
9
#import tensorflow_addons as tfa
10
11
import matplotlib.pyplot as plt
import math
12
from tensorflow.keras.layers import Dense,LeakyReLU,Flatten
13

14
15
16

class Watersheed_Layer(tf.keras.layers.Layer):

Jonathan Juhl's avatar
Jonathan Juhl committed
17
    def __init__(self,mean_matrix,variance_matrix,num_parts):
18
        super(Watersheed_Layer, self).__init__()
Jonathan Juhl's avatar
Jonathan Juhl committed
19
20
21
        self.mean_matrix = mean_matrix
        self.mean_bias = variance_matrix
        self.num_parts = num_parts
22
23

    def build(self,input_shape):
Jonathan Juhl's avatar
Jonathan Juhl committed
24
       
25
        self.kernel = self.add_weight("offset",
Jonathan Juhl's avatar
Jonathan Juhl committed
26
                                  shape=[self.num_parts,input_shape[-1]],trainable=True,initializer=tf.keras.initializers.Constant(self.mean_matrix))
27
     
Jonathan Juhl's avatar
Jonathan Juhl committed
28
29
        self.mean = self.add_weight("offset",
                                  shape=[self.num_parts,input_shape[-1]],trainable=False,initializer=tf.keras.initializers.Constant(self.mean_matrix))
30

Jonathan Juhl's avatar
Jonathan Juhl committed
31
32
33
34
35
36
37
38
39
40
41
42
43

        self.variance = self.add_weight("offset",
                                  shape=[self.num_parts],trainable=False,initializer=tf.keras.initializers.Constant(self.mean_bias))
    
     
    def call(self,noise):
  
        t_noise = tf.tile(tf.expand_dims(noise,axis=1),[1,self.num_parts,1])
        tmp = self.mean*t_noise+tf.expand_dims(tf.expand_dims(self.variance,axis=0),axis=2)
        matrix = tf.matmul(tf.math.l2_normalize(self.kernel,axis=-1),tf.math.l2_normalize(tmp,axis=-1),transpose_b=True)
        s = tf.reduce_sum(1-matrix,axis=(1,2))
      
        return s
Jonathan Juhl's avatar
Jonathan Juhl committed
44
45
    

46

Jonathan Juhl's avatar
Jonathan Juhl committed
47
48
def loss_latent(predict,catagories):
    return tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(predict,catagories))
49

Jonathan Juhl's avatar
Jonathan Juhl committed
50
51
def loss_gen(D_logits_,predict_z,z):

Jonathan Juhl's avatar
Jonathan Juhl committed
52
    g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(tf.ones_like(D_logits_),D_logits_))+loss_latent(predict_z,z)
53
54
    return g_loss

55
56
def loss_encode(d_h1_r,d_h1_f,d_h2_r,d_h2_f,d_h3_r,d_h3_f,d_h4_r,d_h4_f):

Jonathan Juhl's avatar
Jonathan Juhl committed
57
58
        d_h1_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(tf.ones_like(d_h1_r),d_h1_r)) \
                    + tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(tf.zeros_like(d_h1_f),d_h1_f)) 
59

Jonathan Juhl's avatar
Jonathan Juhl committed
60
61
        d_h2_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits( tf.ones_like(d_h2_r),d_h2_r)) \
                    + tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(tf.zeros_like(d_h2_f),d_h2_f))  
62

Jonathan Juhl's avatar
Jonathan Juhl committed
63
64
        d_h3_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(tf.ones_like(d_h3_r),d_h3_r)) \
                    + tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(tf.zeros_like(d_h3_f),d_h3_f))  
65

Jonathan Juhl's avatar
Jonathan Juhl committed
66
67
        d_h4_loss =  tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(tf.ones_like(d_h4_r),d_h4_r)) \
                    + tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(tf.zeros_like(d_h4_f),d_h4_f))  
68
        return d_h1_loss+d_h2_loss+d_h3_loss+d_h4_loss
69

70
      
71

Jonathan Juhl's avatar
Jonathan Juhl committed
72
def loss_disc(D_logits,D_logits_fake,predict_z,z):
73
74
        
        z = tf.cast(z,tf.float32)
Jonathan Juhl's avatar
Jonathan Juhl committed
75
76
        d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(tf.ones_like(D_logits),D_logits))
        d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(tf.zeros_like(D_logits_fake),D_logits_fake))
77
      
78

79
        d_loss = d_loss_real+d_loss_fake
Jonathan Juhl's avatar
Jonathan Juhl committed
80
81
82
83
 
        sub_loss_latent = loss_latent(predict_z,z)
        d_loss += sub_loss_latent
        return d_loss,sub_loss_latent
84
85
86

class transform_3D(tf.keras.layers.Layer):

Jonathan Juhl's avatar
Jonathan Juhl committed
87
    def __init__(self,full_image):
88
        super(transform_3D, self).__init__()
Jonathan Juhl's avatar
Jonathan Juhl committed
89
        self.full_image = full_image
90
91

    def build(self,input_shape):
Jonathan Juhl's avatar
Jonathan Juhl committed
92
93
94
        self.dimensions = input_shape[-2]
        self.channels = input_shape[-1]

95
        x = tf.range(-int(np.floor(self.dimensions/2)),int(np.ceil(self.dimensions/2)))
96
97
98
99
100
101
        X,Y,Z = tf.meshgrid(x,x,x)

        self.coordinates = tf.stack([tf.reshape(X,[-1]), tf.reshape(Y,[-1]),tf.reshape(Z,[-1])],axis=0)
        
        self.kernel = tf.cast(self.coordinates,tf.float32)
     
Jonathan Juhl's avatar
Jonathan Juhl committed
102
    def rotation_map(self,voxels,alpha,beta,gamma,x_translate,y_translate,z_translate):
103
104
105
106

        alpha = 2*3.1415*tf.cast(alpha,tf.float32)
        beta =2*3.1415*tf.cast(beta,tf.float32)
        gamma = 2*3.1415*tf.cast(gamma,tf.float32)
Jonathan Juhl's avatar
Jonathan Juhl committed
107
108
109
        x_translate = tf.cast(x_translate,tf.float32)
        y_translate = tf.cast(y_translate,tf.float32)
        z_translate = tf.cast(z_translate,tf.float32)
110
        #voxels = tf.transpose(voxels,perm=[1,2,3,0,4])
Jonathan Juhl's avatar
Jonathan Juhl committed
111
     
112
        #batchdim = tf.shape(voxels)[-2]        
113
     
114
115
116
        rotation_matrix_x =tf.stack([tf.ones_like(alpha),tf.zeros_like(alpha),tf.zeros_like(alpha),
                                    tf.zeros_like(alpha),tf.cos(alpha), -tf.sin(alpha),
                                    tf.zeros_like(alpha),tf.sin(alpha), tf.cos(alpha)])
117
        rotation_matrix_y = tf.stack([
118
119
120
                                tf.cos(beta),tf.zeros_like(alpha), tf.sin(beta),
                                tf.zeros_like(alpha),tf.ones_like(alpha),tf.zeros_like(alpha),
                                -tf.sin(beta),tf.zeros_like(alpha), tf.cos(beta)])
121
122
123


        rotation_matrix_z = tf.stack([
124
125
126
127
                                    tf.cos(gamma), -tf.sin(gamma),tf.zeros_like(alpha),  
                                    tf.sin(gamma), tf.cos(gamma),tf.zeros_like(alpha),
                                    tf.zeros_like(alpha),tf.zeros_like(alpha),tf.ones_like(alpha)])
        
128
129
130
131
        rotation_matrix_x = tf.reshape(rotation_matrix_x, (3,3))
        rotation_matrix_y = tf.reshape(rotation_matrix_y, (3,3))
        rotation_matrix_z = tf.reshape(rotation_matrix_z, (3,3))
        s = tf.matmul(rotation_matrix_x,tf.matmul(rotation_matrix_y,rotation_matrix_z)) 
Jonathan Juhl's avatar
Jonathan Juhl committed
132
   
133
        r = tf.matmul(tf.matmul(rotation_matrix_x,tf.matmul(rotation_matrix_y,rotation_matrix_z)) ,self.kernel)
134
   
135
        x,y,z = tf.split(r,3,axis=0)
Jonathan Juhl's avatar
Jonathan Juhl committed
136
    
Jonathan Juhl's avatar
Jonathan Juhl committed
137
138
139
        X = tf.reshape(x,[-1])+x_translate*(self.dimensions/self.full_image)
        Y = tf.reshape(y,[-1])+y_translate*(self.dimensions/self.full_image)  
        Z = tf.reshape(z,[-1])+z_translate*(self.dimensions/self.full_image)    
140
141
142
143
144
145
146
147
148
149
150
151
152
        X_lower  = tf.math.floor(X) # tf.clip_by_value(tf.math.floor(X),0,self.dimensions)
        X_upper  = tf.math.ceil(X)  # tf.clip_by_value(tf.math.ceil(X),0,self.dimensions)

        Y_lower  = tf.math.floor(Y)  # tf.clip_by_value(tf.math.floor(Y),0,self.dimensions)
        Y_upper  = tf.math.ceil(Y) # tf.clip_by_value(tf.math.ceil(Y),0,self.dimensions)

        Z_lower  = tf.math.floor(Z)  # tf.clip_by_value(tf.math.floor(Z),0,self.dimensions)
        Z_upper  = tf.math.ceil(Z)  #tf.clip_by_value(tf.math.ceil(Z),0,self.dimensions)

        x_d = (X-X_lower+0.001)/(X_upper-X_lower+0.001)
        y_d = (Y-Y_lower+0.001)/(Y_upper-Y_lower+0.001)
        z_d = (Z-Z_lower+0.001)/(Z_upper-Z_lower+0.001)
        
Jonathan Juhl's avatar
Jonathan Juhl committed
153
154
155
156
157
158
159
160
161
        coord_000 = tf.stack([X_lower,Y_lower,Z_lower],axis=1)+tf.cast(tf.floor(self.dimensions/2),tf.float32)

        coord_001 = tf.stack([X_lower,Y_lower,Z_upper],axis=1)+tf.cast(tf.floor(self.dimensions/2),tf.float32)
        coord_011 = tf.stack([X_lower,Y_upper,Z_upper],axis=1)+tf.cast(tf.floor(self.dimensions/2),tf.float32)
        coord_111 = tf.stack([X_upper,Y_upper,Z_upper],axis=1)+tf.cast(tf.floor(self.dimensions/2),tf.float32)
        coord_101 = tf.stack([X_upper,Y_lower,Z_upper],axis=1)+tf.cast(tf.floor(self.dimensions/2),tf.float32)
        coord_100 = tf.stack([X_upper,Y_lower,Z_lower],axis=1)+tf.cast(tf.floor(self.dimensions/2),tf.float32)
        coord_010 = tf.stack([X_lower,Y_upper,Z_lower],axis=1)+tf.cast(tf.floor(self.dimensions/2),tf.float32)
        coord_110 = tf.stack([X_upper,Y_upper,Z_lower],axis=1)+tf.cast(tf.floor(self.dimensions/2),tf.float32)
162
        #voxels = tf.reshape(voxels,[dimensions**3,channels])
Jonathan Juhl's avatar
Jonathan Juhl committed
163
   
164
        c000 = tf.gather_nd(voxels,tf.cast(coord_000,tf.int32))
Jonathan Juhl's avatar
Jonathan Juhl committed
165

166
       # print(c000);exit()
167
168
169
170
171
172
173
        c001 = tf.gather_nd(voxels,tf.cast(coord_001,tf.int32))
        c011 = tf.gather_nd(voxels,tf.cast(coord_011,tf.int32))
        c111 = tf.gather_nd(voxels,tf.cast(coord_111,tf.int32))
        c101 = tf.gather_nd(voxels,tf.cast(coord_101,tf.int32))
        c100 = tf.gather_nd(voxels,tf.cast(coord_100,tf.int32))
        c010 = tf.gather_nd(voxels,tf.cast(coord_010,tf.int32))
        c110 = tf.gather_nd(voxels,tf.cast(coord_110,tf.int32))
174
175
176
177

        x_d = tf.expand_dims(x_d,axis=1)
        y_d = tf.expand_dims(y_d,axis=1)
        z_d = tf.expand_dims(z_d,axis=1)
178
179
      
        c00 = c000*(1-x_d) + c100*x_d
Jonathan Juhl's avatar
Jonathan Juhl committed
180
     
181
182
183
184
185
186
187
188
189
        c01 = c001*(1-x_d) + c101*x_d
        c10 = c010*(1-x_d) + c110*x_d
        c11 = c011*(1-x_d) + c111*x_d

        c0 = c00*(1-y_d) + c10*y_d
        c1 = c01*(1-y_d) + c11*y_d
    
        c = c0*(1-z_d)+c1*z_d
     
190
        
Jonathan Juhl's avatar
Jonathan Juhl committed
191
192
       
        out = tf.reshape(c,[self.dimensions,self.dimensions,self.dimensions,self.channels])
193
      
194
        return out 
Jonathan Juhl's avatar
Jonathan Juhl committed
195
    def call(self,voxels,alpha,beta,gamma,x_translate,y_translate,z_translate):
196
      
Jonathan Juhl's avatar
Jonathan Juhl committed
197
        image = tf.map_fn(lambda x: self.rotation_map(tf.squeeze(x[0]),x[1],x[2],x[3],x[4],x[5],x[6]),[voxels,alpha,beta,gamma,x_translate,y_translate,z_translate],dtype=tf.float32) 
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
     
        return image

    

def matmul_func(first,second,third):

        a = tf.matmul(first,second,transpose_a=True)#*third
        a = tf.tile(a,[2,1,1])
        third = tf.transpose(third,perm=[2,0,1])
     
        return a*third
def matmul_func_2D(first,second):

        a = tf.matmul(first,second,transpose_a=True)#*third

    
     
        return tf.reshape(tf.reshape(a,[-1]),[2,2])

class Instance_norm(tf.keras.layers.Layer):

    def __init__(self,return_mean):
        super(Instance_norm, self).__init__()
        self.return_mean = return_mean
    
    def build(self,input_shape):
        
        self.offset = self.add_weight("offset",
                                  shape=[input_shape[-1]],trainable=True)
        self.scale  = self.add_weight("scale",
                                  shape=[input_shape[-1]],trainable=True)

    def call(self,x):

Jonathan Juhl's avatar
Jonathan Juhl committed
233
        mean, variance = tf.nn.moments(x, axes=[1, 2])
234
235
        epsilon = 1e-5
        inv = 1/tf.sqrt(variance + epsilon)
Jonathan Juhl's avatar
Jonathan Juhl committed
236
237
        normalized = tf.transpose((tf.transpose(x,perm=[1,2,0,3]) - mean) * inv,perm=[2,0,1,3])
    
238
239
240
241
242
243
244
245
246
247
248
249
250
        if self.return_mean:
            return self.scale * normalized + self.offset, mean, variance
        else:
            return self.scale  * normalized + self.offset

class Spectral_norm(tf.keras.layers.Layer):

    def __init__(self,channels,kernels=3,strides=2):
        super(Spectral_norm, self).__init__()
        self.channels = channels
        self.strides = strides
        self.kernels = kernels
    def build(self,input_shape):
Jonathan Juhl's avatar
Jonathan Juhl committed
251
252
     
        self.w = self.add_weight("kernel",
253
                                  shape=[self.kernels,self.kernels,input_shape[-1],self.channels])
Jonathan Juhl's avatar
Jonathan Juhl committed
254
255
256
257
258
        self.u = self.add_weight("spec_kernel",
                                  shape=[1,self.channels], initializer=tf.keras.initializers.TruncatedNormal(), trainable=False)
        self.bias =  self.add_weight("offset",  
                                  shape=[1,self.channels])
        self.inp_shape = input_shape[-1]
259
260
261

    def call(self,x):

Jonathan Juhl's avatar
Jonathan Juhl committed
262
263
264
265
266
267
268
269
        w = tf.reshape(self.w, [-1, self.channels])
        v_ = tf.matmul(self.u,w,transpose_b=True)
     
        v_hat = tf.math.l2_normalize(v_)
            
        u_ = tf.matmul(v_hat,w)
        
        u_hat = tf.math.l2_normalize(u_)
270
271
272

        sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(u_hat))
        w_norm = w / sigma
Jonathan Juhl's avatar
Jonathan Juhl committed
273
274
        w_norm = tf.reshape(w_norm, [self.kernels,self.kernels,self.inp_shape,self.channels])
        conv = tf.nn.conv2d(x, w_norm, strides=[1, self.strides, self.strides, 1], padding='SAME')+self.bias
275
276
277
278
279
280
        return conv



def draw_from_distribution(width,ctf_params,projections_count):

Jonathan Juhl's avatar
Jonathan Juhl committed
281
282
283
284
285
286
287
    translate_x= np.linspace(-np.floor(width/4),np.ceil(width/4),projections_count)
    outs  = np.random.choice(translate_x,int(0.98*1000))
    translate_x_outer = np.linspace(-np.floor(width/4),np.ceil(width/4),projections_count)
    outs_outer  = np.random.choice(translate_x_outer,int(0.02*1000))
    stdev = np.std(np.concatenate([outs_outer,outs],axis=0))
    translate_x = np.random.normal(size=projections_count,loc=0.0,scale=stdev)
    translate_y = np.random.normal(size=projections_count,loc=0.0,scale=stdev)
Jonathan Juhl's avatar
Jonathan Juhl committed
288
289
    translate_z = np.random.normal(size=projections_count,loc=0.0,scale=stdev)

290
291
292
293
294

    a = np.arange(ctf_params.shape[0])
    sa = np.random.choice(a,projections_count)

    ctf_pr_count = np.take_along_axis(ctf_params,np.expand_dims(sa,axis=1),axis=0)
Jonathan Juhl's avatar
Jonathan Juhl committed
295
  
296
297
    inplane_count  =     np.random.choice(np.random.uniform(0,1,36),projections_count)

Jonathan Juhl's avatar
Jonathan Juhl committed
298
    return ctf_pr_count,translate_x,translate_y,translate_z,inplane_count
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336



def projection(model,gamma,beta,alpha,translations):
    gamma = tf.cast(gamma,tf.float32)
    beta = tf.cast(beta,tf.float32)
    alpha = tf.cast(alpha,tf.float32)
    translations = tf.cast(translations,tf.float32)
    rot_model = tf.transpose(tfa.image.rotate(model,2*3.1415*gamma,interpolation='bilinear'),perm=[1,0,2])
    rot_model = tfa.image.rotate(rot_model,2*3.1415*beta,interpolation='bilinear')

    #projection =tfa.image.rotate(tf.reduce_sum(rot_model,axis=2),2*3.1415*alpha,interpolation='bilinear') # tfa.image.translate(tfa.image.rotate(tf.reduce_sum(rot_model,axis=2),2*3.1415*alpha,interpolation='bilinear'),translations,interpolation='bilinear')
 
    return tf.expand_dims(tf.reduce_sum(rot_model,axis=2),axis=-1)


def projection_map(model,gamma,beta,alpha,translations):
    model = tf.squeeze(tf.cast(model,tf.float32))
    images = tf.map_fn(lambda x: projection(model,x[0],x[1],x[2],x[3]),[gamma,beta,alpha,translations],dtype=tf.float32)
    return images


def translate(image_stack,translate_x,translate_y):

    shifted_images = tfa.image.translate(tf.expand_dims(image_stack,axis=-1),tf.cast(tf.stack([translate_y,translate_x],axis=1),tf.float32))
    return tf.expand_dims(image_stack,axis=-1) #image_stack

def mm_models(model_size,num_linear_params):

    model_vectors = np.random.uniform(low=0.0, high=1.0, size=[model_size**3,num_linear_params])
    q,r = np.linalg.qr(model_vectors)
    return q

def sample_params(batch_size):
    a  = np.random.uniform(size=[batch_size])
    #a = a/np.reduce_sum(a,axis=1)
    return a

337
338
339
340
def ElectronWavelength(kV=300.0):
    # Returns electorn wavelength in Angstroms
    # kV *= 1e3 # ensure Kilovolts for below formula
    # return 12.2639 / np.sqrt( kV + 0.97845 * kV*kV / ( 1e6 ) )
Jonathan Juhl's avatar
Jonathan Juhl committed
341
    return 12.2639 / tf.sqrt(kV * 1e3 + 0.97845 * kV * kV)
342
def apply_ctf(image,ctf_params,KVolts,spherical_abberation,w2):
Jonathan Juhl's avatar
Jonathan Juhl committed
343
    
344
345
    length = tf.shape(image)[2]
    gamma = ElectronWavelength(KVolts)
346
    image = tf.cast(image,tf.float32)
347
    image = tf.transpose(image,perm=[0,3,1,2])
348
349
    ctf_params = tf.cast(ctf_params,tf.float32)
    gamma = tf.cast(gamma,tf.float32)
350
351
    Cs = tf.cast(spherical_abberation,tf.float32)*1e7 
    w2 = tf.cast(w2,tf.float32)
352
353
    scale_dimension = tf.shape(image)[2]
    phase,defocus_1,defocus_2,ast = tf.split(ctf_params,4,axis=1)
354
355
356
357
358
359
    
    phase = tf.cast(phase,tf.float32)
    ast *= -3.1415 / 180.0
    
    w1 = tf.cast(tf.math.sqrt(1 - w2 * w2),tf.float32)

360

Jonathan Juhl's avatar
Jonathan Juhl committed
361
362
    linear_map_1 = tf.concat([[0.0],tf.linspace(1/length,-1/length,length-1)],axis=0)

363

364
365
366
    xmeshtile,ymeshtile = tf.meshgrid(linear_map_1,linear_map_1)
    rmesh = tf.math.sqrt(xmeshtile * xmeshtile + ymeshtile * ymeshtile) / 1.0
    amesh = tf.cast(tf.reshape(tf.math.atan2(ymeshtile,xmeshtile),[-1]),tf.float32)
367

368
    delta_f = 0.5 * (defocus_1 + defocus_2 + (defocus_1 - defocus_2) * tf.math.cos(2.0 * (amesh - ast)))
Jonathan Juhl's avatar
Jonathan Juhl committed
369
 
370
371
372
    rmesh2 = tf.cast(rmesh*rmesh,tf.float32)
    rmesh2 = tf.reshape(rmesh2,[-1])
    amesh = tf.reshape(amesh,[-1])
Jonathan Juhl's avatar
Jonathan Juhl committed
373

374
    xi  = tf.constant(3.1415,dtype=tf.float32) * gamma * rmesh2 * (delta_f - tf.constant(0.5,dtype=tf.float32) * gamma**2 * rmesh2 * Cs)
Jonathan Juhl's avatar
Jonathan Juhl committed
375
376

    #print(np.amin(xi.numpy()),np.amax(xi.numpy()));exit()
377
378
    #print(g);exit()
    #print(xi);exit()
Jonathan Juhl's avatar
Jonathan Juhl committed
379
    CTF = -w1 * tf.math.sin(xi) - w2 * tf.math.cos(xi)
380
381
    s = tf.cast(tf.sqrt(tf.cast(tf.shape(CTF)[1],tf.float32)),tf.int32)
    CTF = tf.cast(tf.reshape(CTF,[tf.shape(CTF)[0],s,s]),tf.complex64)
Jonathan Juhl's avatar
Jonathan Juhl committed
382
   
383
384
    image = tf.cast(image,tf.complex64)

Jonathan Juhl's avatar
Jonathan Juhl committed
385
386
387
388
    ctf_image  = tf.cast(tf.signal.ifft2d(tf.signal.fft2d(tf.squeeze(image))*(-CTF)),tf.float32)
  
   
    return ctf_image
Jonathan Juhl's avatar
Jonathan Juhl committed
389
390
391
392
393
394
395
396
397
398
399
"""
data = np.load('data.npy')

data = np.reshape(data,(1,256,256,256,1))
t = transform_3D(256)
v = np.asarray([0.0])
k = np.asarray([128])
out = t(data,v,v,v,k,k,v)
plt.imshow(np.squeeze(np.sum(out,axis=1)))
plt.savefig('try_it.png')
#import mrcfile"""
Jonathan Juhl's avatar
Jonathan Juhl committed
400
"""V = 300
401
lambdas  = 10**(-4)*10**(-6)*12.25*10**(-10)/np.sqrt(V*10**3)
402

403
404
405
spher_abb =  2.700000
ac = 0.100000 
ctf_params = np.asarray([[0.0,19196.580078, 18916.789062, 57.389999 ]])
406
407


Jonathan Juhl's avatar
Jonathan Juhl committed
408
409
410
411
412
413
414
415
416
417
418
419
420
421
with mrcfile.open('/emcc/misser11/EMPIAR_10317/without_ctf.mrcs') as mrc:
    voxel = mrc.data
    t = 0
    for i in range(voxel.shape[0]):
        print(t)
       # print(np.squeeze(voxel[t]).shape);exit()
        plt.imshow(np.squeeze(voxel[t]))
        plt.savefig('with_out_ctf_%i.png' %t)
        plt.clf()
        t+=1

#voxel = apply_ctf(voxel,ctf_params,V,spher_abb,ac)
#plt.imshow(np.squeeze(voxel))
#plt.savefig('without_ctf.png')
422

Jonathan Juhl's avatar
Jonathan Juhl committed
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
with mrcfile.open('/emcc/misser11/EMPIAR_10317/with_ctf.mrcs') as mrc:
   voxel = mrc.data

   t = 0
   for i in range(voxel.shape[0]):
        print(t)
        plt.imshow(np.squeeze(voxel[t]))
        plt.savefig('with_ctf_%i.png' %t)
        plt.clf()
        t+=1
with mrcfile.open('/emcc/misser11/EMPIAR_10317/out_noise.mrcs') as mrc:
   voxel = mrc.data

   t = 0
   for i in range(voxel.shape[0]):
        print(t)
        plt.imshow(np.squeeze(voxel[t]))
        plt.savefig('with_noise_%i.png' %t)
        plt.clf()
        t+=1

Jonathan Juhl's avatar
Jonathan Juhl committed
444
"""